dapa46/task-6a02e23da6fe2e4ac16acf65
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(needs_fixes): 1 исправлений, 0 отстояно — main.py

Browse Source
This commit is contained in:
Danil Parunin 5f1b81b8-4f5d-11e8-9c2d-fa7ae01bbebc
2026-06-16 16:20:09 +00:00
parent ca9c3783f6
commit 6e7697b740
1 changed files with 101 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files
main.py
+101 -145
View File
@@ -1,35 +1,29 @@
"""
# main.py – RAG‑agent with Qdrant, OpenRouter, and LangChain
# -----------------------------------------------------------------
# This script implements a simple RAG agent that can search and add
# documents to a Qdrant vector store. The agent is built with
# LangChain's `create_agent` and uses OpenRouter for both the LLM and
# embeddings. The code follows the "Исправить" section of the
# assignment and includes detailed comments explaining design choices.
# -----------------------------------------------------------------
DeepAgents RAG Agent with Qdrant and Ollama embeddings
"""
import os
import asyncio
import argparse
from pathlib import Path
from typing import List
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from langchain_community.document_loaders import TextLoader
from langchain_community.document_loaders import DirectoryLoader
from langchain_community.vectorstores import Qdrant
from langchain_ollama.embeddings import OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.agents import create_agent, AgentExecutor, AgentType
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
# -----------------------------------------------------------------
# Configuration – all secrets are read from environment variables.
# -----------------------------------------------------------------
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
raise RuntimeError("OPENAI_API_KEY environment variable is required")
# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
# Environment variables
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "") # required for OpenRouter LLM
OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
# LLM – OpenRouter gpt-oss-20b:free (free tier)
# LLM via OpenRouter
llm = ChatOpenAI(
model="openai/gpt-oss-20b:free",
base_url="https://openrouter.ai/api/v1",
@@ -37,150 +31,112 @@ llm = ChatOpenAI(
temperature=0.0,
)
# Embeddings – OpenAI text-embedding-3-small via OpenRouter
embeddings = OpenAIEmbeddings(
model="text-embedding-3-small",
base_url="https://openrouter.ai/api/v1",
api_key=OPENAI_API_KEY,
)
# Embeddings via Ollama
embeddings = OllamaEmbeddings(model="nomic-embed-text", base_url=OLLAMA_BASE_URL)
# Qdrant client – assumes a local Qdrant instance running on default port
qdrant_url = os.getenv("QDRANT_URL", "http://localhost:6333")
vector_store = Qdrant(
client=None, # will be created lazily by Qdrant wrapper
collection_name="knowledge",
# Qdrant vector store
qdrant_vector_store = QdrantVectorStore(
client_kwargs={"url": QDRANT_URL},
collection_name="knowledge_base",
embeddings=embeddings,
url=qdrant_url,
# Create collection if not exists
create_collection=True,
# Use cosine similarity
distance="cosine",
)
# -----------------------------------------------------------------
# Tool definitions – these are the only tools the agent can use.
# -----------------------------------------------------------------
# Text splitter
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------
@tool
def search_knowledge_base(query: str, max_results: int = 3) -> str:
"""Search the knowledge base for relevant information.
Parameters
----------
query: str
The search query.
max_results: int, optional
Number of top results to return (default 3).
"""
docs = vector_store.similarity_search(query, k=max_results)
"""Semantic search in the knowledge base."""
docs = qdrant_vector_store.similarity_search(query, k=max_results)
if not docs:
return "No results found."
return "\n\n---\n\n".join([f"{doc.metadata.get('title', 'Untitled')}\n{doc.page_content}" for doc in docs])
return "\n\n".join(f"{i+1}. {doc.page_content}" for i, doc in enumerate(docs))
@tool
def add_to_knowledge_base(content: str, title: str = "Untitled") -> str:
"""Add a new document (or chunk) to the knowledge base.
def add_to_knowledge_base(content: str, title: str = "document") -> str:
"""Add content to the knowledge base after chunking."""
# Split content into chunks
chunks = text_splitter.split_text(content)
metadatas = [{"title": title, "chunk_index": i} for i in range(len(chunks))]
# Add to Qdrant
qdrant_vector_store.add_texts(chunks, metadatas=metadatas)
return f"Added {len(chunks)} chunks from '{title}'."
Parameters
----------
content: str
The text content to add.
title: str, optional
A human‑readable title for the document.
"""
doc = {
"page_content": content,
"metadata": {"title": title},
}
vector_store.add_documents([doc])
return f"Added document '{title}'."
# ---------------------------------------------------------------------------
# Backend setup
# ---------------------------------------------------------------------------
backend = CompositeBackend([
LocalShellBackend(workspace_dir="./workspace"),
FilesystemBackend(),
])
# -----------------------------------------------------------------
# Agent setup – using LangChain's create_agent with a custom system prompt.
# -----------------------------------------------------------------
SYSTEM_PROMPT = (
"You are a helpful assistant with access to a knowledge base. "
"Use the tools `search_knowledge_base` and `add_to_knowledge_base` "
"to answer user queries. If the user asks to add information, "
"use `add_to_knowledge_base`. If the user asks for information, "
"use `search_knowledge_base`. Do not fabricate facts."
)
agent = create_agent(
llm=llm,
# ---------------------------------------------------------------------------
# Agent creation
# ---------------------------------------------------------------------------
agent = create_deep_agent(
model=llm,
tools=[search_knowledge_base, add_to_knowledge_base],
system_prompt=SYSTEM_PROMPT,
agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
backend=backend,
system_prompt="You are a helpful assistant with access to a knowledge base. Use the provided tools to search and add information.",
)
executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True)
# ---------------------------------------------------------------------------
# Document loader for initialization
# ---------------------------------------------------------------------------
async def load_documents_from_dir(directory: str):
"""Load all .txt files from a directory into the knowledge base."""
dir_path = Path(directory)
for file_path in dir_path.rglob("*.txt"):
content = file_path.read_text(encoding="utf-8")
title = file_path.stem
await agent.ainvoke(
{"messages": [HumanMessage(content=f"/add {content}")], "metadata": {"title": title}},
{"configurable": {"thread_id": "init-session"}},
)
# -----------------------------------------------------------------
# Document ingestion – split into chunks and add to Qdrant.
# -----------------------------------------------------------------
def ingest_directory(directory: str, chunk_size: int = 1000, chunk_overlap: int = 200):
"""Load all text files from *directory*, split into chunks, and store.
Parameters
----------
directory: str
Path to the directory containing documents.
chunk_size: int, optional
Size of each chunk in characters.
chunk_overlap: int, optional
Overlap between consecutive chunks.
"""
loader = DirectoryLoader(directory, glob="**/*.txt")
documents = loader.load()
splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
chunks = splitter.split_documents(documents)
# Convert LangChain Document objects to dicts expected by Qdrant
docs_to_add = []
for doc in chunks:
title = doc.metadata.get("source", "Untitled")
docs_to_add.append({
"page_content": doc.page_content,
"metadata": {"title": title, "source": doc.metadata.get("source", "")},
})
vector_store.add_documents(docs_to_add)
print(f"Ingested {len(docs_to_add)} chunks into the knowledge base.")
# -----------------------------------------------------------------
# CLI – simple interactive loop.
# -----------------------------------------------------------------
async def main():
parser = argparse.ArgumentParser(description="RAG Agent CLI")
parser.add_argument("--ingest", type=str, help="Path to directory to ingest")
args = parser.parse_args()
if args.ingest:
ingest_directory(args.ingest)
return
print("RAG Agent ready. Type /quit to exit.")
# ---------------------------------------------------------------------------
# CLI client
# ---------------------------------------------------------------------------
async def cli():
print("DeepAgents RAG CLI. Commands: /add <text>, /search <query>, /quit")
while True:
user_input = input("You: ")
user_input = input("> ")
if user_input.strip() == "/quit":
print("Goodbye!")
break
if user_input.startswith("/add "):
# Expected format: /add <title> | <content>
try:
_, rest = user_input.split("/add ", 1)
title, content = rest.split("|", 1)
title = title.strip()
content = content.strip()
result = await executor.ainvoke({"messages": [HumanMessage(content=f"Add document {title}")], "configurable": {"thread_id": "session-1"}})
# Directly call tool to add content
add_to_knowledge_base(content, title)
print("Agent: Document added.")
except Exception as e:
print(f"Error parsing /add command: {e}")
continue
if user_input.startswith("/search "):
query = user_input[len("/search "):].strip()
result = await executor.ainvoke({"messages": [HumanMessage(content=f"Search for {query}")], "configurable": {"thread_id": "session-1"}})
print("Agent:", result["messages"][-1].content)
continue
# Default: normal chat
result = await executor.ainvoke({"messages": [HumanMessage(content=user_input)], "configurable": {"thread_id": "session-1"}})
print("Agent:", result["messages"][-1].content)
content = user_input[5:].strip()
response = await agent.ainvoke(
{"messages": [HumanMessage(content=f"/add {content}")], "metadata": {"title": "user_input"}},
{"configurable": {"thread_id": "cli-session"}},
)
print(response["messages"][-1].content)
elif user_input.startswith("/search "):
query = user_input[8:].strip()
response = await agent.ainvoke(
{"messages": [HumanMessage(content=f"/search {query}")], "metadata": {"query": query}},
{"configurable": {"thread_id": "cli-session"}},
)
print(response["messages"][-1].content)
else:
print("Unknown command. Use /add, /search, or /quit.")
# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------
async def main():
# Optional: load initial documents from a folder named 'data'
data_dir = Path("./data")
if data_dir.exists():
await load_documents_from_dir(str(data_dir))
await cli()
if __name__ == "__main__":
asyncio.run(main())
"""