From fe5d3887448f73903d5f003d20eedf0b0ef0a936 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=94=D0=B0=D0=BD=D0=B8=D0=B8=D0=BB=20=D0=92=D0=B8=D0=BA?=
 =?UTF-8?q?=D1=82=D0=BE=D1=80=D0=BE=D0=B2?= <mr.face2002@gmail.com>
Date: Tue, 30 Jun 2026 17:35:42 +0000
Subject: [PATCH] =?UTF-8?q?add:=20main.py=20=E2=80=94=20=D0=AD=D0=BA=D0=B7?=
 =?UTF-8?q?=D0=B0=D0=BC=D0=B5=D0=BD:=20RAG-=D0=B0=D0=B3=D0=B5=D0=BD=D1=82?=
 =?UTF-8?q?=20=D1=81=20ChromaDB=20=D0=B8=20=D0=B2=D0=B5=D0=B1-=D0=BF=D0=BE?=
 =?UTF-8?q?=D0=B8=D1=81=D0=BA=D0=BE=D0=BC?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 main.py | 136 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 136 insertions(+)
 create mode 100644 main.py

diff --git a/main.py b/main.py
new file mode 100644
index 0000000..145da66
--- /dev/null
+++ b/main.py
@@ -0,0 +1,136 @@
+import os
+import asyncio
+from pathlib import Path
+
+from dotenv import load_dotenv
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain_chroma import Chroma
+from langchain_core.documents import Document
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.tools import tool
+from langchain_community.tools.tavily_search import TavilySearchResults
+
+from deepagents import create_deep_agent
+from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
+from langchain_core.messages import HumanMessage
+
+load_dotenv()
+
+# ---------- 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,
+)
+
+# ---------- Vector Store ----------
+embeddings = OpenAIEmbeddings(
+    model="text-embedding-3-small",
+    base_url="https://openrouter.ai/api/v1",
+    api_key=os.getenv("OPENAI_API_KEY"),
+)
+
+def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
+    return Chroma(
+        collection_name="knowledge",
+        embedding_function=embeddings,
+        persist_directory=persist_directory,
+    )
+
+vectorstore = create_vectorstore()
+
+def load_documents(directory: str, vectorstore: Chroma) -> None:
+    path = Path(directory)
+    if not path.is_dir():
+        return
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    docs = []
+    for file_path in path.rglob("*"):
+        if file_path.suffix.lower() not in {".txt", ".md"}:
+            continue
+        text = file_path.read_text(encoding="utf-8")
+        chunks = splitter.split_text(text)
+        for i, chunk in enumerate(chunks):
+            metadata = {"source": str(file_path), "chunk_index": i}
+            docs.append(Document(page_content=chunk, metadata=metadata))
+    if docs:
+        vectorstore.add_documents(docs)
+
+# Load initial documents (run once or each start)
+load_documents("./documents", vectorstore)
+
+# ---------- Tools ----------
+@tool
+def search_local_kb(query: str, top_k: int = 3) -> str:
+    """
+    Perform a semantic search in the local Chroma knowledge base.
+    Returns the concatenated contents of the most relevant documents.
+    """
+    docs = vectorstore.similarity_search(query, k=top_k)
+    if not docs:
+        return "No relevant information found in the local knowledge base."
+    return "\n---\n".join(doc.page_content for doc in docs)
+
+tavily_tool = TavilySearchResults(max_results=5)
+
+@tool
+def web_search(query: str) -> str:
+    """
+    Search the web using Tavily and return a short summary of the top results.
+    """
+    results = tavily_tool.run(query)
+    if not results:
+        return "No web results found."
+    # results is a list of dicts; extract title and url
+    lines = []
+    for r in results:
+        title = r.get("title", "No title")
+        url = r.get("url", "")
+        lines.append(f"{title}: {url}")
+    return "\n".join(lines)
+
+# ---------- Backend ----------
+backend = CompositeBackend(
+    [
+        LocalShellBackend(workspace_dir="./workspace"),
+        FilesystemBackend(),
+    ]
+)
+
+# ---------- Agent ----------
+system_prompt = (
+    "You are an AI assistant that can answer questions using two sources: "
+    "a local knowledge base (accessed via the tool `search_local_kb`) and the web (accessed via the tool `web_search`). "
+    "Decide which tool to use based on the user query. "
+    "When you use a tool, include the source name in your final answer: "
+    "`Source: chromadb` for local knowledge, `Source: tavily` for web results. "
+    "If both sources are needed, combine them and list both sources."
+)
+
+agent = create_deep_agent(
+    model=llm,
+    tools=[search_local_kb, web_search],
+    backend=backend,
+    system_prompt=system_prompt,
+)
+
+# ---------- CLI ----------
+async def chat_loop() -> None:
+    print("AI Assistant (type 'exit' to quit)")
+    thread_id = "session-1"
+    while True:
+        user_input = input("\nYou: ").strip()
+        if user_input.lower() in {"exit", "quit"}:
+            print("Goodbye!")
+            break
+
+        result = await agent.ainvoke(
+            {"messages": [HumanMessage(content=user_input)]},
+            {"configurable": {"thread_id": thread_id}},
+        )
+        answer = result["messages"][-1].content
+        print(f"\nAssistant: {answer}")
+
+if __name__ == "__main__":
+    asyncio.run(chat_loop())
\ No newline at end of file