feat: solution for 6a1864f78a94f887e50d46da

solutions/6a1864f78a94f887e50d46da/solution.py

@@ -0,0 +1,87 @@
+# -------------------- 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"):
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
+    vector_store = Chroma(
+        collection_name="rag_collection",
+        embedding_function=embeddings,
+        persist_directory=persist_directory,
+    )
+    return vector_store
+
+def load_documents(directory: str, vectorstore):
+    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")
+        docs.extend(splitter.create_documents([text]))
+    vectorstore.add_documents(docs)
+
+# -------------------- tools.py --------------------
+from langchain.tools import tool
+from langchain_ollama import ChatOllama
+
+@tool
+def search_local_kb(query: str, top_k: int = 3) -> str:
+    """Semantic search in the local ChromaDB knowledge base."""
+    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
+    docs = retriever.invoke({"query": query})["documents"]
+    return "\n".join(doc.page_content for doc in docs)
+
+@tool
+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})
+    return "\n".join(f"{r['title']}: {r['url']}" for r in results)
+
+# -------------------- agent.py --------------------
+from langchain.agents import create_agent
+from langchain_ollama import ChatOllama
+
+llm = ChatOllama(model="llama3", temperature=0.2)
+
+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,
+)
+
+# -------------------- main.py --------------------
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+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()
+
+    print("Chat started. Type 'exit' to quit.")
+    while True:
+        user_input = input("\nЗапрос: ").strip()
+        if 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}")