diff --git a/main.py b/main.py index 89056a5..d04b61e 100644 --- a/main.py +++ b/main.py @@ -1,93 +1,109 @@ import os -import sys +import asyncio from pathlib import Path -from dotenv import load_dotenv - 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.tools import tool -from langchain.agents import AgentExecutor, create_openai_tools_agent -from langchain_community.tools.tavily import TavilySearchResults +from langchain_tavily import TavilySearchResults +from deepagents import create_deep_agent +from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend -# Load env variables -load_dotenv() - -# ---------- LLM and embeddings ---------- +# ---------- LLM ---------- llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", api_key=os.getenv("OPENAI_API_KEY"), temperature=0.0, ) -embeddings = OllamaEmbeddings(model="nomic-embed-text") -# ---------- Vectorstore ---------- +# ---------- Backend ---------- +backend = CompositeBackend([ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), +]) + +# ---------- Vectorstore utilities ---------- PERSIST_DIR = Path("./chroma_db") PERSIST_DIR.mkdir(parents=True, exist_ok=True) -vectorstore = Chroma(persist_directory=str(PERSIST_DIR), embedding_function=embeddings) -# Load documents if collection empty -if not vectorstore.get_collection().count(): - docs_dir = Path("./documents") - if docs_dir.exists(): - splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) - for file in docs_dir.glob("**/*.*"): - if file.suffix.lower() in {".txt", ".md"}: - text = file.read_text(encoding="utf-8") - chunks = splitter.split_text(text) - vectorstore.add_texts(chunks, ids=[f"{file.name}_{i}" for i in range(len(chunks))]) +# Create or load Chroma vectorstore +vectorstore = Chroma( + persist_directory=str(PERSIST_DIR), + embedding_function=OllamaEmbeddings(model="nomic-embed-text"), +) + +# 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.extend(splitter.split_text(text)) + # Convert to LangChain Documents + 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("search_local_kb") +@tool def search_local_kb(query: str, top_k: int = 3) -> str: - """Semantic search in local ChromaDB knowledge base.""" + """Semantic search in the local ChromaDB knowledge base.""" retriever = vectorstore.as_retriever(search_kwargs={"k": top_k}) docs = retriever.get_relevant_documents(query) if not docs: return "No relevant documents found in local KB." return "\n---\n".join(doc.page_content for doc in docs) -@tool("web_search") +@tool def web_search(query: str) -> str: """Web search using Tavily.""" - tavily = TavilySearchResults(max_results=3) + tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY")) 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) # ---------- Agent ---------- -SYSTEM_PROMPT = ( - "You are an AI assistant. For questions about local documents, use the tool 'search_local_kb'. " - "For up-to-date information or news, use 'web_search'. Always indicate the source 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." + ), ) -tools = [search_local_kb, web_search] -agent = create_openai_tools_agent(llm, tools, system_message=SYSTEM_PROMPT) -agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True) - # ---------- CLI ---------- - -def main(): +async def main(): print("RAG Agent with ChromaDB and Tavily. Type 'exit' to quit.") while True: - try: - query = input("\nЗапрос: ") - except (EOFError, KeyboardInterrupt): - print("\nBye!") + user_input = input("\nЗапрос: ") + if user_input.lower() in {"exit", "quit"}: + print("Goodbye!") break - if query.strip().lower() in {"exit", "quit"}: - print("Bye!") - break - result = agent_executor.invoke({"input": query}) - # The tool name is stored in the tool_calls field of the result - tool_name = result.get("tool_calls", [{}])[0].get("name", "unknown") - source = "tavily" if tool_name == "web_search" else "chromadb" - print(f"\n[{'Web Search' if source=='tavily' else 'Local KB'}] {result.get('output', '')}\n") - print(f"Источник: {source}\n") + 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 + reply = result["messages"][-1].content + print(f"\n{reply}") if __name__ == "__main__": - main() + asyncio.run(main())