diff --git a/main.py b/main.py index fb7f61b..2dd5ee5 100644 --- a/main.py +++ b/main.py @@ -1,13 +1,11 @@ import os import asyncio -from pathlib import Path from dotenv import load_dotenv - -from langchain_openai import ChatOpenAI, OpenAIEmbeddings +from langchain_ollama import ChatOllama +from langchain_ollama import OllamaEmbeddings from langchain_chroma import Chroma from langchain_core.documents import Document from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain_tavily import TavilySearchResults from langchain.tools import tool from deepagents import create_deep_agent from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend @@ -15,97 +13,92 @@ from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeB # Load environment variables load_dotenv() -# ---------- LLM and Embeddings ---------- -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 = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), -) - # ---------- Vector Store ---------- -CHROMA_DIR = Path("./chroma_db") -vector_store = Chroma( - collection_name="knowledge", - embedding_function=embeddings, - persist_directory=str(CHROMA_DIR), -) + +def create_vectorstore(persist_directory="./chroma_db"): + """Create a Chroma vector store with Ollama embeddings.""" + embeddings = OllamaEmbeddings(model="nomic-embed-text") + return Chroma(persist_directory=persist_directory, embedding_function=embeddings) + + +def load_documents(directory, vectorstore): + """Load .txt and .md files from *directory*, chunk them, and add to *vectorstore*. + The function preserves the file name in metadata for later reference. + """ + splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + 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() + docs = splitter.split_text(text) + documents = [Document(page_content=chunk, metadata={"source": fname}) for chunk in docs] + vectorstore.add_documents(documents) # ---------- Tools ---------- + +vectorstore = create_vectorstore() + @tool def search_local_kb(query: str, top_k: int = 3) -> str: - """Semantic search in the local knowledge base.""" - docs = vector_store.similarity_search(query, k=top_k) + """Semantic search in the local knowledge base (ChromaDB).""" + 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(f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs)) + return "No relevant local knowledge found." + return "\n---\n".join([f"{doc.metadata.get('source', 'unknown')}\n{doc.page_content}" for doc in docs]) @tool def web_search(query: str) -> str: """Web search using Tavily.""" - tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY")) - results = tavily.run(query) + from tavily import TavilyClient + client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY")) + results = client.search(query, max_results=3) if not results: return "No web results found." - return "\n---\n".join(f"{i+1}. {res['title']}\n{res['content'][:200]}..." for i, res in enumerate(results)) + return "\n---\n".join([f"{r['title']}\n{r['content']}" for r in results]) + +# ---------- Agent ---------- + +llm = ChatOllama(model="llama3", temperature=0.0) -# ---------- Backend ---------- backend = CompositeBackend([ LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend(), ]) -# ---------- Agent ---------- +system_prompt = ( + "You are an AI assistant with access to two tools: " + "search_local_kb for local knowledge and web_search for up-to-date information. " + "When answering a user query, first decide which tool is appropriate. " + "If the answer can be derived from the local documents, use search_local_kb; " + "otherwise use web_search. " + "Always indicate the source of the information in your response: " + "[Local KB] or [Web Search]." +) + agent = create_deep_agent( model=llm, tools=[search_local_kb, web_search], backend=backend, - system_prompt=( - "You are a helpful assistant. For any user query, decide whether to use the local knowledge base or perform a web search. " - "If the query is about recent events, news, or requires up‑to‑date information, use the web_search tool. " - "Otherwise, use search_local_kb. " - "Always return the source used in the response (either 'chromadb' or 'tavily')." - ), + system_prompt=system_prompt, ) -# ---------- Document Loader ---------- -def load_documents(directory: str, vectorstore: Chroma): - """Load .txt and .md files from a directory, chunk them, and add to the vector store.""" - splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) - docs = [] - for file_path in Path(directory).glob("**/*"): - if file_path.suffix.lower() in {".txt", ".md"}: - text = file_path.read_text(encoding="utf-8") - chunks = splitter.split_text(text) - docs.extend([Document(page_content=chunk, metadata={"source": str(file_path)}) for chunk in chunks]) - if docs: - vectorstore.add_documents(docs) - vectorstore.persist() - -# ---------- CLI ---------- async def main(): - # Ensure vector store is loaded - if not CHROMA_DIR.exists() or not any(CHROMA_DIR.iterdir()): - print("Loading documents into ChromaDB…") - load_documents("./documents", vector_store) - print("RAG agent ready. Type 'exit' to quit.") + 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": [{"role": "user", "content": user_input}]}, {"configurable": {"thread_id": "session-1"}}, ) - # Extract last message content - content = result["messages"][-1].content - print(f"\nОтвет:\n{content}") + # The last message is the assistant's reply + reply = result["messages"][-1].content + print(f"\n{reply}") if __name__ == "__main__": asyncio.run(main())