code quality fixes

Author: heliamoh

src/agent/profiles/base.py

@@ -52,20 +52,12 @@ class BaseState(InputState, OutputState, total=False):
 class BaseGraphBuilder:
     """Base class for all graph builders with common preprocessing and postprocessing."""
 
-    def __init__(
-        self, 
-        llm: BaseChatModel, 
-        embedding: Embeddings
-        ) -> None:
+    def __init__(self, llm: BaseChatModel, embedding: Embeddings) -> None:
         """Initialize with LLM and embedding models."""
         self.preprocessing_workflow: Runnable = create_preprocessing_workflow(llm)
         self.search_workflow: Runnable = create_search_workflow(llm)
 
-    async def preprocess(
-        self, 
-        state: BaseState, 
-        config: RunnableConfig
-        ) -> BaseState:
+    async def preprocess(self, state: BaseState, config: RunnableConfig) -> BaseState:
         """Run the complete preprocessing workflow and map results to state."""
         result: PreprocessingState = await self.preprocessing_workflow.ainvoke(
             PreprocessingState(
@@ -77,10 +69,7 @@ async def preprocess(
 
         return self._map_preprocessing_result(result)
 
-    def _map_preprocessing_result(
-        self, 
-        result: PreprocessingState
-        ) -> BaseState:
+    def _map_preprocessing_result(self, result: PreprocessingState) -> BaseState:
         """Map preprocessing results to BaseState with defaults."""
         return BaseState(
             rephrased_input=result["rephrased_input"],
@@ -90,11 +79,7 @@ def _map_preprocessing_result(
             detected_language=result.get("detected_language", DEFAULT_LANGUAGE),
         )
 
-    async def postprocess(
-        self, 
-        state: BaseState, 
-        config: RunnableConfig
-        ) -> BaseState:
+    async def postprocess(self, state: BaseState, config: RunnableConfig) -> BaseState:
         """Postprocess that preserves existing state and conditionally adds search results."""
         search_results: list[WebSearchResult] = []
 
@@ -113,4 +98,9 @@ async def postprocess(
             search_results = result["search_results"]
 
         # Create new state with updated additional_content
-        return BaseState(**{**state, "additional_content": AdditionalContent(search_results=search_results)})
+        return BaseState(
+            **{
+                **state,
+                "additional_content": AdditionalContent(search_results=search_results),
+            }
+        )

src/agent/profiles/react_to_me.py

@@ -23,11 +23,7 @@ class ReactToMeState(BaseState):
 class ReactToMeGraphBuilder(BaseGraphBuilder):
     """Graph builder for ReactToMe profile with Reactome-specific functionality."""
 
-    def __init__(
-        self, 
-        llm: BaseChatModel, 
-        embedding: Embeddings
-        ) -> None:
+    def __init__(self, llm: BaseChatModel, embedding: Embeddings) -> None:
         """Initialize ReactToMe graph builder with required components."""
         super().__init__(llm, embedding)
 
@@ -69,25 +65,20 @@ def _build_workflow(self) -> StateGraph:
         return state_graph
 
     async def preprocess(
-        self, 
-        state: ReactToMeState, 
-        config: RunnableConfig
+        self, state: ReactToMeState, config: RunnableConfig
     ) -> ReactToMeState:
         """Run preprocessing workflow."""
         result = await super().preprocess(state, config)
         return ReactToMeState(**result)
 
     async def proceed_with_research(
-        self, 
-        state: ReactToMeState
+        self, state: ReactToMeState
     ) -> Literal["Continue", "Finish"]:
         """Determine whether to proceed with research based on safety check."""
         return "Continue" if state["safety"] == SAFETY_SAFE else "Finish"
 
     async def generate_unsafe_response(
-        self, 
-        state: ReactToMeState, 
-        config: RunnableConfig
+        self, state: ReactToMeState, config: RunnableConfig
     ) -> ReactToMeState:
         """Generate appropriate refusal response for unsafe queries."""
         final_answer_message = await self.unsafe_answer_generator.ainvoke(
@@ -120,9 +111,7 @@ async def generate_unsafe_response(
         )
 
     async def call_model(
-        self, 
-        state: ReactToMeState, 
-        config: RunnableConfig
+        self, state: ReactToMeState, config: RunnableConfig
     ) -> ReactToMeState:
         """Generate response using Reactome RAG for safe queries."""
         result: dict[str, Any] = await self.reactome_rag.ainvoke(
@@ -147,9 +136,6 @@ async def call_model(
         )
 
 
-def create_reactome_graph(
-    llm: BaseChatModel, 
-    embedding: Embeddings
-    ) -> StateGraph:
+def create_reactome_graph(llm: BaseChatModel, embedding: Embeddings) -> StateGraph:
     """Create and return the ReactToMe workflow graph."""
     return ReactToMeGraphBuilder(llm, embedding).uncompiled_graph

src/retrievers/csv_chroma.py

@@ -77,11 +77,7 @@ def list_chroma_subdirectories(directory: Path) -> List[str]:
 class HybridRetriever:
     """Advanced hybrid retriever supporting RRF, parallel processing, and multi-source search."""
 
-    def __init__(
-        self, 
-        embedding: Embeddings, 
-        embeddings_directory: Path
-        ):
+    def __init__(self, embedding: Embeddings, embeddings_directory: Path):
 
         self.embedding = embedding
         self.embeddings_directory = embeddings_directory
@@ -159,25 +155,17 @@ def _create_vector_retriever(self, subdirectory: str) -> Optional[object]:
             return None
 
     async def _search_with_bm25(
-        self, 
-        query: str, 
-        retriever: BM25Retriever
+        self, query: str, retriever: BM25Retriever
     ) -> List[Document]:
         """Search using BM25 retriever asynchronously."""
         return await asyncio.to_thread(retriever.get_relevant_documents, query)
 
-    async def _search_with_vector(
-        self, 
-        query: str, 
-        retriever: Any
-    ) -> List[Document]:
+    async def _search_with_vector(self, query: str, retriever: Any) -> List[Document]:
         """Search using vector retriever asynchronously."""
         return await asyncio.to_thread(retriever.get_relevant_documents, query)
 
     async def _execute_hybrid_search(
-        self, 
-        query: str, 
-        subdirectory: str
+        self, query: str, subdirectory: str
     ) -> List[Document]:
         """Execute hybrid search (BM25 + vector) for a single query on a subdirectory."""
         retriever_info = self._retrievers.get(subdirectory)
@@ -223,9 +211,7 @@ def _generate_document_identifier(self, document: Document) -> str:
         return hashlib.md5(document.page_content.encode()).hexdigest()
 
     async def _apply_reciprocal_rank_fusion(
-        self, 
-        queries: List[str], 
-        subdirectory: str
+        self, queries: List[str], subdirectory: str
     ) -> List[Document]:
         """Apply Reciprocal Rank Fusion to results from multiple queries on a subdirectory."""
         logger.info(
@@ -325,8 +311,7 @@ async def ainvoke(self, inputs: Dict[str, Any]) -> str:
 
 
 def create_hybrid_retriever(
-    embedding: Embeddings, 
-    embeddings_directory: Path
+    embedding: Embeddings, embeddings_directory: Path
 ) -> HybridRetriever:
     """Create a hybrid retriever with RRF and parallel processing support."""
     try: