diff --git a/solutions/6a1864f78a94f887e50d46da/solution.py b/solutions/6a1864f78a94f887e50d46da/solution.py index 710af4e..c94fecd 100644 --- a/solutions/6a1864f78a94f887e50d46da/solution.py +++ b/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}") \ No newline at end of file + for msg in result["messages"]: + # `msg` is a Pydantic model; use the `.content` attribute + if hasattr(msg, "content"): + print(msg.content) \ No newline at end of file