187 lines
7.3 KiB
Python
187 lines
7.3 KiB
Python
"""
|
||
# 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.
|
||
# -----------------------------------------------------------------
|
||
|
||
import os
|
||
import asyncio
|
||
import argparse
|
||
from pathlib import Path
|
||
|
||
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
|
||
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_text_splitters import RecursiveCharacterTextSplitter
|
||
from langchain.agents import create_agent, AgentExecutor, AgentType
|
||
|
||
# -----------------------------------------------------------------
|
||
# 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")
|
||
|
||
# LLM – OpenRouter gpt-oss-20b:free (free tier)
|
||
llm = ChatOpenAI(
|
||
model="openai/gpt-oss-20b:free",
|
||
base_url="https://openrouter.ai/api/v1",
|
||
api_key=OPENAI_API_KEY,
|
||
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,
|
||
)
|
||
|
||
# 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",
|
||
embeddings=embeddings,
|
||
url=qdrant_url,
|
||
)
|
||
|
||
# -----------------------------------------------------------------
|
||
# Tool definitions – these are the only tools the agent can use.
|
||
# -----------------------------------------------------------------
|
||
@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)
|
||
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])
|
||
|
||
@tool
|
||
def add_to_knowledge_base(content: str, title: str = "Untitled") -> str:
|
||
"""Add a new document (or chunk) to the knowledge base.
|
||
|
||
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}'."
|
||
|
||
# -----------------------------------------------------------------
|
||
# 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,
|
||
tools=[search_knowledge_base, add_to_knowledge_base],
|
||
system_prompt=SYSTEM_PROMPT,
|
||
agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
|
||
)
|
||
|
||
executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True)
|
||
|
||
# -----------------------------------------------------------------
|
||
# 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.")
|
||
while True:
|
||
user_input = input("You: ")
|
||
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)
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|
||
"""
|