feat: solution for 6a1864f78a94f887e50d46da

solutions/6a1864f78a94f887e50d46da/solution.py

@@ -1,87 +1,112 @@
-# -------------------- vectorstore.py --------------------
+# vectorstore.py
 from pathlib import Path
+
 from langchain_chroma import Chroma
 from langchain_ollama import OllamaEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 
-def create_vectorstore(persist_directory: str = "./chroma_db"):
+
+def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
+    """Create a persistent Chroma vector store with Ollama embeddings."""
     embeddings = OllamaEmbeddings(model="nomic-embed-text")
-    vector_store = Chroma(
+    return Chroma(
         collection_name="rag_collection",
         embedding_function=embeddings,
         persist_directory=persist_directory,
     )
-    return vector_store
 
-def load_documents(directory: str, vectorstore):
+
+def load_documents(directory: str, vectorstore: Chroma) -> None:
+    """Load .txt and .md files from `directory`, chunk them, and add to the store."""
     splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
     docs = []
-    for file_path in Path(directory).glob("*.txt"):
-        text = file_path.read_text(encoding="utf-8")
-        docs.extend(splitter.create_documents([text]))
-    for file_path in Path(directory).glob("*.md"):
-        text = file_path.read_text(encoding="utf-8")
+    for file in Path(directory).glob("*.txt") | Path(directory).glob("*.md"):
+        text = file.read_text(encoding="utf-8")
         docs.extend(splitter.create_documents([text]))
     vectorstore.add_documents(docs)
 
-# -------------------- tools.py --------------------
-from langchain.tools import tool
-from langchain_ollama import ChatOllama
 
-@tool
+# tools.py
+from langchain.tools import tool
+
+from langchain_chroma import Chroma
+from langchain_ollama import OllamaEmbeddings
+
+
+@tool("search_local_kb")
 def search_local_kb(query: str, top_k: int = 3) -> str:
-    """Semantic search in the local ChromaDB knowledge base."""
+    """Semantic search in local knowledge base."""
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
+    vectorstore = Chroma(
+        collection_name="rag_collection",
+        embedding_function=embeddings,
+        persist_directory="./chroma_db",
+    )
     retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
-    docs = retriever.invoke({"query": query})["documents"]
+    docs = retriever.invoke(query)
     return "\n".join(doc.page_content for doc in docs)
 
-@tool
+
+@tool("web_search")
 def web_search(query: str) -> str:
-    """Web search using Tavily."""
-    from langchain_tavily import TavilySearchResults
-    tavily = TavilySearchResults(api_key=__import__("os").environ["TAVILY_API_KEY"])
-    results = tavily.invoke({"query": query})
+    """Search the web using Tavily."""
+    from langchain_tavily import TavilyAPIWrapper
+
+    tavily = TavilyAPIWrapper()
+    results = tavily.run(query)
     return "\n".join(f"{r['title']}: {r['url']}" for r in results)
 
-# -------------------- agent.py --------------------
+
+# agent.py
 from langchain.agents import create_agent
 from langchain_ollama import ChatOllama
 
-llm = ChatOllama(model="llama3", temperature=0.2)
+from tools import search_local_kb, web_search
 
-system_prompt = """
-You are an assistant that answers user questions.
-If the answer can be found in the local knowledge base, use `search_local_kb`.
-Otherwise, use `web_search`.  
-Always indicate the source of your answer: either "chromadb" or "tavily".
-"""
 
-agent = create_agent(
-    model=llm,
-    tools=[search_local_kb, web_search],
-    system_prompt=system_prompt,
-)
+def create_rag_agent():
+    llm = ChatOllama(model="llama3", temperature=0.2)
+    tools = [search_local_kb, web_search]
+    system_prompt = (
+        "You are an assistant that answers user questions.\n"
+        "If the answer can be found in local documents, use search_local_kb.\n"
+        "Otherwise, use web_search. Return the answer followed by a line\n"
+        "\"Source: chromadb\" or \"Source: tavily\"."
+    )
+    return create_agent(model=llm, tools=tools, system_prompt=system_prompt)
 
-# -------------------- main.py --------------------
+
+# main.py
 import os
+from pathlib import Path
+
 from dotenv import load_dotenv
 
-load_dotenv()
+from vectorstore import create_vectorstore, load_documents
+from agent import create_rag_agent
+
+
+def init_db():
+    """Create and populate the Chroma DB if it does not exist."""
+    db_path = Path("./chroma_db")
+    if not db_path.exists():
+        store = create_vectorstore()
+        load_documents("documents", store)
+        store.persist()
+
 
 if __name__ == "__main__":
-    # Initialize vectorstore and load documents if not already loaded
-    vectorstore = create_vectorstore()
-    if not vectorstore.get_collection().count():
-        load_documents("documents", vectorstore)
-        vectorstore.persist()
+    load_dotenv()  # Loads TAVILY_API_KEY and any other env vars
+    init_db()
+    agent = create_rag_agent()
 
-    print("Chat started. Type 'exit' to quit.")
     while True:
-        user_input = input("\nЗапрос: ").strip()
-        if user_input.lower() in ("exit", "quit", "выход"):
+        user_input = input("Запрос: ").strip()
+        if not user_input or user_input.lower() in ("exit", "quit", "выход"):
             break
+
         result = agent.invoke({"messages": [{"role": "human", "content": user_input}]})
-        ai_msg = result["messages"][-1]
-        print(f"[{ai_msg.tool_calls[0]['name'].capitalize()}] {ai_msg.content}")
-        source = "chromadb" if ai_msg.tool_calls[0]["name"] == "search_local_kb" else "tavily"
-        print(f"Источник: {source}")
+        for msg in result["messages"]:
+            # `msg` is a Pydantic model; use the `.content` attribute
+            if hasattr(msg, "content"):
+                print(msg.content)