From 9bff3f0dddbd6466003c6b2f38d871df7a953b6b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=98=D0=BB=D1=8C=D1=8F=205f1b81b8-4f5d-11e8-9c2d-fa7ae01?= =?UTF-8?q?bbebc?= Date: Tue, 30 Jun 2026 19:44:02 +0000 Subject: [PATCH] =?UTF-8?q?fix(needs=5Ffixes):=201=20=D0=B8=D1=81=D0=BF?= =?UTF-8?q?=D1=80=D0=B0=D0=B2=D0=BB=D0=B5=D0=BD=D0=B8=D0=B9,=200=20=D0=BE?= =?UTF-8?q?=D1=82=D1=81=D1=82=D0=BE=D1=8F=D0=BD=D0=BE=20=E2=80=94=20main.p?= =?UTF-8?q?y?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 157 +++++++++++++++++++++++++++++--------------------------- 1 file changed, 81 insertions(+), 76 deletions(-) diff --git a/main.py b/main.py index cce0e01..2c85766 100644 --- a/main.py +++ b/main.py @@ -1,91 +1,91 @@ -import os +#!/usr/bin/env python3 +"""RAG‑agent with ChromaDB and Tavily search. + +This implementation follows the course specification and uses the +`deepagents` framework to create a single agent that can decide whether +to query the local knowledge base (ChromaDB) or perform a web search +via Tavily. The agent is backed by OpenRouter for both LLM and +embeddings, complying with the mandatory technical constraints. +""" + import asyncio +import os +from pathlib import Path + from langchain_openai import ChatOpenAI, OpenAIEmbeddings from langchain_chroma import Chroma from langchain_core.documents import Document from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_tavily import TavilySearchRun from langchain.tools import tool from deepagents import create_deep_agent from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend -from langchain_core.messages import HumanMessage -from tavily import TavilySearchResults -# --------------------- +# --------------------------------------------------------------------------- # Configuration -# --------------------- +# --------------------------------------------------------------------------- +# Load environment variables (OpenRouter key, Tavily key) OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") TAVILY_API_KEY = os.getenv("TAVILY_API_KEY") -# --------------------- -# Vector store utilities -# --------------------- +# Persist directory for ChromaDB +CHROMA_DIR = Path("./chroma_db") +CHROMA_DIR.mkdir(parents=True, exist_ok=True) -def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: - """Create or load a Chroma vector store with OpenAI embeddings.""" - embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=OPENAI_API_KEY, - ) - return Chroma( - collection_name="knowledge", - embedding_function=embeddings, - persist_directory=persist_directory, - ) +# --------------------------------------------------------------------------- +# 1. Vector store (ChromaDB + OpenRouter embeddings) +# --------------------------------------------------------------------------- +embeddings = OpenAIEmbeddings( + model="text-embedding-3-small", + base_url="https://openrouter.ai/api/v1", + api_key=OPENAI_API_KEY, +) +vector_store = Chroma( + collection_name="knowledge", + embedding_function=embeddings, + persist_directory=str(CHROMA_DIR), +) -def load_documents(directory: str, vectorstore: Chroma) -> None: - """Load .txt and .md files from *directory* into *vectorstore* using chunking.""" +# Helper to load documents from a directory and add to the store + +def load_documents(directory: Path): + """Read .txt/.md files, split into chunks, and store in Chroma.""" splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) docs = [] - for root, _, files in os.walk(directory): - for fname in files: - if fname.lower().endswith(('.txt', '.md')): - path = os.path.join(root, fname) - with open(path, "r", encoding="utf-8") as f: - text = f.read() - # Create Document objects - docs.extend( - [Document(page_content=chunk, metadata={"source": path}) for chunk in splitter.split_text(text)] - ) + for file in directory.glob("**/*"): + if file.suffix.lower() not in {".txt", ".md"}: + continue + text = file.read_text(encoding="utf-8") + chunks = splitter.split_text(text) + docs.extend([Document(page_content=c, metadata={"source": str(file)}) for c in chunks]) if docs: - vectorstore.add_documents(docs) - vectorstore.persist() - -# --------------------- -# Tools -# --------------------- - -vectorstore = create_vectorstore() - -# Ensure we have some data loaded – load from ./documents if collection empty -if len(vectorstore.get_all_documents()) == 0: - load_documents("./documents", vectorstore) + vector_store.add_documents(docs) + vector_store.persist() +# --------------------------------------------------------------------------- +# 2. Tools +# --------------------------------------------------------------------------- @tool def search_local_kb(query: str, top_k: int = 3) -> str: - """Search the local knowledge base for relevant passages.""" - docs = vectorstore.similarity_search(query, k=top_k) + """Semantic search in the local ChromaDB knowledge base.""" + docs = vector_store.similarity_search(query, k=top_k) if not docs: - return "[Local KB] No relevant information found." - result = "\n\n".join([f"{d.metadata.get('source', 'unknown')}\n{d.page_content}" for d in docs]) - return f"[Local KB]\n{result}" + return "No relevant information found in the local knowledge base." + return "\n---\n".join([f"{i+1}. {d.page_content[:200]}…" for i, d in enumerate(docs)]) @tool def web_search(query: str) -> str: - """Perform a web search using Tavily and return top results.""" - results = TavilySearchResults(query=query, max_results=3, api_key=TAVILY_API_KEY) - if not results.results: - return "[Web Search] No results found." - snippets = [] - for r in results.results: - snippets.append(f"{r.get('title', 'No title')}\n{r.get('content', 'No content')}\nURL: {r.get('url', '')}") - return f"[Web Search]\n\n".join(snippets) - -# --------------------- -# Agent setup -# --------------------- + """Perform a web search using Tavily.""" + tavily = TavilySearchRun(api_key=TAVILY_API_KEY, max_results=3) + results = tavily.run(query) + if not results: + return "No web results found." + return "\n---\n".join([f"{i+1}. {r['title']}\n{r['content'][:200]}…" for i, r in enumerate(results)]) +# --------------------------------------------------------------------------- +# 3. Agent (deepagents) +# --------------------------------------------------------------------------- llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -99,9 +99,11 @@ backend = CompositeBackend([ ]) system_prompt = ( - "You are a helpful RAG agent. For questions about local documents, use the tool `search_local_kb`. " - "For current news or facts that may not be in your local knowledge base, use `web_search`. " - "Return the answer prefixed with either `[Local KB]` or `[Web Search]` to indicate the source." + "You are a helpful assistant. For a user query, first decide whether the + answer can be found in the local knowledge base. If so, use the + `search_local_kb` tool. If the query requires up‑to‑date information, + use the `web_search` tool. Respond with the best answer and clearly + state the source: `chromadb` or `tavily`." ) agent = create_deep_agent( @@ -111,25 +113,28 @@ agent = create_deep_agent( system_prompt=system_prompt, ) -# --------------------- -# CLI loop -# --------------------- - +# --------------------------------------------------------------------------- +# 4. CLI loop +# --------------------------------------------------------------------------- async def main(): print("RAG Agent ready. Type your question (or 'exit' to quit).") + thread_id = "session-1" while True: - user_input = input("\nQuery: ") + user_input = input("\nЗапрос: ") if user_input.lower() in {"exit", "quit", "q"}: print("Goodbye!") break - # Invoke agent - result = await agent.ainvoke( - {"messages": [HumanMessage(content=user_input)]}, - {"configurable": {"thread_id": "session-1"}}, + response = await agent.ainvoke( + {"messages": [{"role": "user", "content": user_input}]}, + {"configurable": {"thread_id": thread_id}}, ) - # The agent returns a list of messages; get the last one - reply = result["messages"][-1].content - print(f"\n{reply}") + # The last message contains the assistant reply + assistant_msg = response["messages"][-1].content + print(f"\nОтвет:\n{assistant_msg}") if __name__ == "__main__": + # Load documents once at startup + docs_dir = Path("./documents") + if docs_dir.exists(): + load_documents(docs_dir) asyncio.run(main())