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

This commit is contained in:
+113 -92
View File
@@ -1,29 +1,38 @@
"""RAG Agent with Qdrant and OpenRouter.
This script implements the assignment requirements:
* Two tools `search_knowledge_base` and `add_to_knowledge_base` are defined with the `@tool` decorator.
* A Qdrant vector store is used for semantic search. Documents are split into chunks with a
`RecursiveCharacterTextSplitter` that has `chunk_overlap=100` as requested.
* The agent is created with LangChains `create_agent` (the "Исправить" instruction overrides the
earlier requirement to use `create_deep_agent`).
* A simple CLI allows adding documents, searching the knowledge base and quitting.
The code is selfcontained and can be run directly after installing the dependencies listed in
`requirements.txt`.
""" """
DeepAgents RAG Agent with Qdrant and Ollama embeddings
"""
import os import os
import asyncio import asyncio
from pathlib import Path import pathlib
from typing import List from typing import List
from langchain_openai import ChatOpenAI from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.messages import HumanMessage from langchain_core.messages import HumanMessage
from langchain.tools import tool from langchain.tools import tool
from langchain_ollama.embeddings import OllamaEmbeddings from langchain_community.vectorstores import Qdrant
from langchain_qdrant import QdrantVectorStore from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_text_splitters import RecursiveCharacterTextSplitter
from deepagents import create_deep_agent from langchain.agents import create_agent, AgentExecutor, AgentType
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Configuration # Configuration
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Environment variables OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "") # required for OpenRouter LLM if not OPENAI_API_KEY:
OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434") raise RuntimeError("OPENAI_API_KEY environment variable is required")
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
# LLM via OpenRouter # LLM and embeddings via OpenRouter
llm = ChatOpenAI( llm = ChatOpenAI(
model="openai/gpt-oss-20b:free", model="openai/gpt-oss-20b:free",
base_url="https://openrouter.ai/api/v1", base_url="https://openrouter.ai/api/v1",
@@ -31,112 +40,124 @@ llm = ChatOpenAI(
temperature=0.0, temperature=0.0,
) )
# Embeddings via Ollama embeddings = OpenAIEmbeddings(
embeddings = OllamaEmbeddings(model="nomic-embed-text", base_url=OLLAMA_BASE_URL) model="text-embedding-3-small",
base_url="https://openrouter.ai/api/v1",
# Qdrant vector store api_key=OPENAI_API_KEY,
qdrant_vector_store = QdrantVectorStore(
client_kwargs={"url": QDRANT_URL},
collection_name="knowledge_base",
embeddings=embeddings,
# Create collection if not exists
create_collection=True,
# Use cosine similarity
distance="cosine",
) )
# ---------------------------------------------------------------------------
# Vector store setup (Qdrant)
# ---------------------------------------------------------------------------
# Qdrant is expected to be running locally on the default port 6333.
# If you need a different host/port, adjust the `url` parameter.
vector_store = Qdrant.from_existing_index(
collection_name="knowledge",
embeddings=embeddings,
url="http://localhost:6333",
)
# ---------------------------------------------------------------------------
# Text splitter # Text splitter
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) # ---------------------------------------------------------------------------
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=1000,
chunk_overlap=100, # as required by the assignment
)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Tools # Tools
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@tool @tool
def search_knowledge_base(query: str, max_results: int = 3) -> str: def search_knowledge_base(query: str, max_results: int = 3) -> str:
"""Semantic search in the knowledge base.""" """Search the knowledge base for relevant information.
docs = qdrant_vector_store.similarity_search(query, k=max_results)
Parameters
----------
query: str
The search query.
max_results: int, optional
Number of top results to return. Defaults to 3.
"""
docs = vector_store.similarity_search(query, k=max_results)
if not docs: if not docs:
return "No results found." return "No results found."
return "\n\n".join(f"{i+1}. {doc.page_content}" for i, doc in enumerate(docs)) return "\n\n---\n\n".join(doc.page_content for doc in docs)
@tool @tool
def add_to_knowledge_base(content: str, title: str = "document") -> str: def add_to_knowledge_base(content: str, title: str = "document") -> str:
"""Add content to the knowledge base after chunking.""" """Add content to the knowledge base.
# Split content into chunks
Parameters
----------
content: str
The raw text to add.
title: str, optional
A title for the document. Defaults to "document".
"""
# Split into chunks and create Document objects
chunks = text_splitter.split_text(content) chunks = text_splitter.split_text(content)
metadatas = [{"title": title, "chunk_index": i} for i in range(len(chunks))] from langchain_core.documents import Document
# Add to Qdrant
qdrant_vector_store.add_texts(chunks, metadatas=metadatas) docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks]
return f"Added {len(chunks)} chunks from '{title}'." vector_store.add_documents(docs)
return f"Added {len(docs)} chunks for title '{title}'."
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Backend setup # Agent creation (LangChain create_agent)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
backend = CompositeBackend([ # The system prompt instructs the agent to use the knowledge base tools.
LocalShellBackend(workspace_dir="./workspace"), SYSTEM_PROMPT = (
FilesystemBackend(), "You are an AI 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, store it. If the user asks for information, search the base."
# ---------------------------------------------------------------------------
# Agent creation
# ---------------------------------------------------------------------------
agent = create_deep_agent(
model=llm,
tools=[search_knowledge_base, add_to_knowledge_base],
backend=backend,
system_prompt="You are a helpful assistant with access to a knowledge base. Use the provided tools to search and add information.",
) )
# --------------------------------------------------------------------------- agent = create_agent(
# Document loader for initialization llm=llm,
# --------------------------------------------------------------------------- tools=[search_knowledge_base, add_to_knowledge_base],
async def load_documents_from_dir(directory: str): system_prompt=SYSTEM_PROMPT,
"""Load all .txt files from a directory into the knowledge base.""" agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
dir_path = Path(directory) )
for file_path in dir_path.rglob("*.txt"):
content = file_path.read_text(encoding="utf-8") agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base])
title = file_path.stem
await agent.ainvoke(
{"messages": [HumanMessage(content=f"/add {content}")], "metadata": {"title": title}},
{"configurable": {"thread_id": "init-session"}},
)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# CLI client # CLI
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
async def cli(): async def handle_user_input(user_input: str) -> str:
print("DeepAgents RAG CLI. Commands: /add <text>, /search <query>, /quit") 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 = add_to_knowledge_base(content, title)
return result
except ValueError:
return "Invalid format. Use: /add <title> | <content>"
elif user_input.startswith("/search "):
query = user_input[len("/search "):].strip()
return search_knowledge_base(query)
elif user_input == "/quit":
return "quit"
else:
# Forward to the agent
response = await agent_executor.ainvoke({"messages": [HumanMessage(content=user_input)]})
return response["messages"][-1].content
async def main():
print("RAG Agent CLI. Commands: /add <title> | <content>, /search <query>, /quit")
while True: while True:
user_input = input("> ") user_input = input("> ")
if user_input.strip() == "/quit": if not user_input:
continue
result = await handle_user_input(user_input)
if result == "quit":
print("Goodbye!") print("Goodbye!")
break break
if user_input.startswith("/add "): print(result)
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__": if __name__ == "__main__":
asyncio.run(main()) asyncio.run(main())