fix: code fix.

src/agent_hackathon/create_vector_db.py

@@ -138,12 +138,12 @@ class VectorDBCreator:
         logger.info("Pipeline finished.")
 
 
-if __name__ == "__main__":
-    logger.info("Script started.")
-    # Optionally load environment variables if needed
-    _ = load_dotenv(dotenv_path=find_dotenv(raise_error_if_not_found=True))
-    creator = VectorDBCreator(
-        data_path=f"{PROJECT_ROOT_DIR}/data/cs_data_arxiv.json", db_uri="arxiv_docs.db"
-    )
-    creator.run()
-    logger.info("Script finished.")
+# if __name__ == "__main__":
+#     logger.info("Script started.")
+#     # Optionally load environment variables if needed
+#     _ = load_dotenv(dotenv_path=find_dotenv(raise_error_if_not_found=True))
+#     creator = VectorDBCreator(
+#         data_path=f"{PROJECT_ROOT_DIR}/data/cs_data_arxiv.json", db_uri="arxiv_docs.db"
+#     )
+#     creator.run()
+#     logger.info("Script finished.")

src/agent_hackathon/generate_arxiv_responses.py

@@ -21,6 +21,7 @@ class ArxivResponseGenerator:
         """Initializes the ArxivResponseGenerator."""
         self.vector_store_path = vector_store_path
         self.client = self._initialise_client()
+        self.retriever = self._initialise_retriever()
         logger.info("ArxivResponseGenerator initialized.")
 
     def _initialise_retriever(self) -> Any:
@@ -40,7 +41,7 @@ class ArxivResponseGenerator:
         )
         retriever = retriever_class.build_retriever_engine()
         logger.info("Retriever engine initialized.")
-        return retriever, retriever_class
+        return retriever
 
     def _initialise_client(self) -> InferenceClient:
         """
@@ -68,11 +69,15 @@ class ArxivResponseGenerator:
             str: Formatted response from the LLM.
         """
         logger.info(f"Retrieving arXiv papers for query: {query}")
-        retriever, retriever_class = self._initialise_retriever()
-        retrieved_content = json.dumps(
-            obj=[(i.get_content(), i.metadata) for i in retriever.retrieve(query)]
-        )
-        logger.info("Retrieved content from vector DB.")
+
+        try:
+            retrieved_content = json.dumps(
+                obj=[(i.get_content(), i.metadata) for i in self.retriever.retrieve(query)]
+            )
+            logger.info("Retrieved content from vector DB.")
+        except Exception as err:
+            logger.error(f"Error retrieving from vector DB: {err}")
+            raise
         completion = self.client.chat.completions.create(
             model="meta-llama/Llama-4-Scout-17B-16E-Instruct",
             temperature=0.1,
@@ -89,17 +94,15 @@ class ArxivResponseGenerator:
             ],
         )
         logger.info("Received completion from LLM.")
-        retriever_class.vector_store.client.close()
-        logger.info("Closed vector store client.")
         return completion.choices[0].message.content
 
 
-if __name__ == "__main__":
-    logger.info("Script started.")
-    generator = ArxivResponseGenerator(
-        vector_store_path=PROJECT_ROOT_DIR / "db/arxiv_docs.db"
-    )
-    query = "deep learning for NLP"  # Example query, replace as needed
-    result = generator.retrieve_arxiv_papers(query=query)
-    print(result)
-    logger.info("Script finished.")
+# if __name__ == "__main__":
+#     logger.info("Script started.")
+#     generator = ArxivResponseGenerator(
+#         vector_store_path=PROJECT_ROOT_DIR / "db/arxiv_docs.db"
+#     )
+#     query = "deep learning for NLP"  # Example query, replace as needed
+#     result = generator.retrieve_arxiv_papers(query=query)
+#     print(result)
+#     logger.info("Script finished.")

src/agent_hackathon/query_vector_db.py

@@ -5,7 +5,6 @@ from dotenv import find_dotenv, load_dotenv
 from huggingface_hub import login
 from llama_index.core import VectorStoreIndex
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
-from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
 from llama_index.vector_stores.milvus import MilvusVectorStore
 
 from src.agent_hackathon.consts import PROJECT_ROOT_DIR
@@ -23,7 +22,6 @@ class RetrieverEngineBuilder:
         self,
         hf_token_env: str = "HF_TOKEN",
         embedding_model: str = "Qwen/Qwen3-Embedding-0.6B",
-        llm_model: str = "meta-llama/Llama-4-Scout-17B-16E-Instruct",
         vector_store: MilvusVectorStore = None,
         device: str = "cpu",
     ) -> None:
@@ -33,27 +31,21 @@ class RetrieverEngineBuilder:
         Args:
             hf_token_env: Environment variable name for HuggingFace token.
             embedding_model: Name of the embedding model.
-            llm_model: Name of the LLM model.
             vector_store: An instance of MilvusVectorStore.
             device: Device to run the embedding model on.
         """
         self.hf_token_env = hf_token_env
         self.embedding_model = embedding_model
-        self.llm_model = llm_model
         self.vector_store = vector_store
         self.device = device
 
         logger.info("Initializing RetrieverEngineBuilder.")
-        self._login_huggingface()
-        self._load_env()
+        # self._login_huggingface()
+        # self._load_env()
 
         self.embed_model = HuggingFaceEmbedding(
             model_name=self.embedding_model, device=self.device
         )
-        self.llm = HuggingFaceInferenceAPI(
-            model=self.llm_model,
-            provider="auto",
-        )
         logger.info("RetrieverEngineBuilder initialized.")
 
     def _login_huggingface(self) -> None:
@@ -65,7 +57,7 @@ class RetrieverEngineBuilder:
     def _load_env(self) -> None:
         """Load environment variables from .env file."""
         logger.info("Loading environment variables.")
-        _ = load_dotenv(dotenv_path=find_dotenv(raise_error_if_not_found=True))
+        _ = load_dotenv(dotenv_path=find_dotenv(raise_error_if_not_found=False))
         logger.info("Environment variables loaded.")
 
     def build_retriever_engine(self) -> Any: