From 7b972b123ac23280d135ce083573f08eaa4c3e8f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=93=D0=BB=D0=B5=D0=B1=20=D0=9D=D0=B8=D0=BA=D0=B8=D1=88?= =?UTF-8?q?=D0=B8=D0=BD?= Date: Thu, 4 Jun 2026 15:55:29 +0000 Subject: [PATCH] add main.py --- main.py | 79 +++++++++++++++++++++++++++++---------------------------- 1 file changed, 40 insertions(+), 39 deletions(-) diff --git a/main.py b/main.py index d04b61e..3dd71ce 100644 --- a/main.py +++ b/main.py @@ -4,12 +4,13 @@ from pathlib import Path from langchain_openai import ChatOpenAI from langchain_core.messages import HumanMessage from langchain.tools import tool -from langchain_ollama import OllamaEmbeddings -from langchain_chroma import Chroma -from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain_tavily import TavilySearchResults from deepagents import create_deep_agent from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend +from langchain_chroma import Chroma +from langchain_ollama import OllamaEmbeddings +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_ollama import OllamaEmbeddings +from langchain_tavily import TavilySearchResults # ---------- LLM ---------- llm = ChatOpenAI( @@ -27,42 +28,35 @@ backend = CompositeBackend([ # ---------- Vectorstore utilities ---------- PERSIST_DIR = Path("./chroma_db") -PERSIST_DIR.mkdir(parents=True, exist_ok=True) -# Create or load Chroma vectorstore -vectorstore = Chroma( - persist_directory=str(PERSIST_DIR), - embedding_function=OllamaEmbeddings(model="nomic-embed-text"), -) +def create_vectorstore(persist_directory: str = "./chroma_db"): + embeddings = OllamaEmbeddings(model="nomic-embed-text") + return Chroma(persist_directory=persist_directory, embedding_function=embeddings) -# Load documents from a directory into the vectorstore def load_documents(directory: str, vectorstore): - docs = [] splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) - for file_path in Path(directory).glob("**/*"): - if file_path.suffix.lower() in {".txt", ".md"}: - text = file_path.read_text(encoding="utf-8") + docs = [] + for file in Path(directory).glob("**/*"): + if file.suffix.lower() in {".txt", ".md"}: + text = file.read_text(encoding="utf-8") docs.extend(splitter.split_text(text)) - # Convert to LangChain Documents + # Convert to Document objects from langchain_core.documents import Document documents = [Document(page_content=chunk) for chunk in docs] vectorstore.add_documents(documents) vectorstore.persist() -# Load documents once at startup (if not already loaded) -if not any(PERSIST_DIR.iterdir()): - load_documents("./documents", vectorstore) - # ---------- Tools ---------- @tool def search_local_kb(query: str, top_k: int = 3) -> str: """Semantic search in the local ChromaDB knowledge base.""" + vectorstore = create_vectorstore(PERSIST_DIR) retriever = vectorstore.as_retriever(search_kwargs={"k": top_k}) - docs = retriever.get_relevant_documents(query) + docs = retriever.invoke(query) if not docs: - return "No relevant documents found in local KB." - return "\n---\n".join(doc.page_content for doc in docs) + return "No relevant local documents found." + return "\n---\n".join(doc.page_content for doc in docs) + "\n[Source: chromadb]" @tool def web_search(query: str) -> str: @@ -71,39 +65,46 @@ def web_search(query: str) -> str: results = tavily.run(query) if not results: return "No web results found." - return "\n---\n".join(f"{r['title']}\n{r['url']}\n{r.get('content', '')}" for r in results) + snippets = [f"{r['title']}\n{r['content']}" for r in results] + return "\n---\n".join(snippets) + "\n[Source: tavily]" # ---------- Agent ---------- +SYSTEM_PROMPT = ( + "You are an AI assistant that can answer questions using either a local knowledge base or the web. " + "If the question refers to documents in the local folder, use the `search_local_kb` tool. " + "If the question is about recent events or requires up‑to‑date information, use the `web_search` tool. " + "Always include the source tag (`[Source: chromadb]` or `[Source: tavily]`) in your answer." +) + agent = create_deep_agent( model=llm, tools=[search_local_kb, web_search], backend=backend, - system_prompt=( - "You are an AI assistant that answers user questions.\n" - "If the question is about information that should be in the local knowledge base,\n" - "use the search_local_kb tool.\n" - "If the question requires up‑to‑date information from the web,\n" - "use the web_search tool.\n" - "Always indicate the source of the answer in the format:\n" - "[Source: chromadb] or [Source: tavily] before the answer." - ), + system_prompt=SYSTEM_PROMPT, ) # ---------- CLI ---------- -async def main(): - print("RAG Agent with ChromaDB and Tavily. Type 'exit' to quit.") +async def chat_loop(): + print("RAG Agent ready. Type 'exit' to quit.") while True: user_input = input("\nЗапрос: ") - if user_input.lower() in {"exit", "quit"}: + if user_input.strip().lower() == "exit": print("Goodbye!") break result = await agent.ainvoke( {"messages": [HumanMessage(content=user_input)]}, {"configurable": {"thread_id": "session-1"}}, ) - # The agent returns a list of messages; the last is the assistant reply + # The last message is the assistant's reply reply = result["messages"][-1].content - print(f"\n{reply}") + print(reply) +# ---------- Initialization ---------- if __name__ == "__main__": - asyncio.run(main()) + # Ensure vectorstore exists and load documents if empty + vectorstore = create_vectorstore(PERSIST_DIR) + if not vectorstore.get_all_documents(): + print("Loading documents into ChromaDB...") + load_documents("./documents", vectorstore) + print("Documents loaded.") + asyncio.run(chat_loop())