164 lines
6.0 KiB
Python
164 lines
6.0 KiB
Python
"""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 LangChain’s `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 self‑contained and can be run directly after installing the dependencies listed in
|
||
`requirements.txt`.
|
||
"""
|
||
|
||
import os
|
||
import asyncio
|
||
import pathlib
|
||
from typing import List
|
||
|
||
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
|
||
from langchain_core.messages import HumanMessage
|
||
from langchain.tools import tool
|
||
from langchain_community.vectorstores import Qdrant
|
||
from langchain_community.document_loaders import TextLoader
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
from langchain.agents import create_agent, AgentExecutor, AgentType
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Configuration
|
||
# ---------------------------------------------------------------------------
|
||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||
if not OPENAI_API_KEY:
|
||
raise RuntimeError("OPENAI_API_KEY environment variable is required")
|
||
|
||
# LLM and embeddings via OpenRouter
|
||
llm = ChatOpenAI(
|
||
model="openai/gpt-oss-20b:free",
|
||
base_url="https://openrouter.ai/api/v1",
|
||
api_key=OPENAI_API_KEY,
|
||
temperature=0.0,
|
||
)
|
||
|
||
embeddings = OpenAIEmbeddings(
|
||
model="text-embedding-3-small",
|
||
base_url="https://openrouter.ai/api/v1",
|
||
api_key=OPENAI_API_KEY,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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 = RecursiveCharacterTextSplitter(
|
||
chunk_size=1000,
|
||
chunk_overlap=100, # as required by the assignment
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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. Defaults to 3.
|
||
"""
|
||
docs = vector_store.similarity_search(query, k=max_results)
|
||
if not docs:
|
||
return "No results found."
|
||
return "\n\n---\n\n".join(doc.page_content for doc in docs)
|
||
|
||
@tool
|
||
def add_to_knowledge_base(content: str, title: str = "document") -> str:
|
||
"""Add content to the knowledge base.
|
||
|
||
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)
|
||
from langchain_core.documents import Document
|
||
|
||
docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks]
|
||
vector_store.add_documents(docs)
|
||
return f"Added {len(docs)} chunks for title '{title}'."
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Agent creation (LangChain create_agent)
|
||
# ---------------------------------------------------------------------------
|
||
# The system prompt instructs the agent to use the knowledge base tools.
|
||
SYSTEM_PROMPT = (
|
||
"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 = create_agent(
|
||
llm=llm,
|
||
tools=[search_knowledge_base, add_to_knowledge_base],
|
||
system_prompt=SYSTEM_PROMPT,
|
||
agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
|
||
)
|
||
|
||
agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base])
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI
|
||
# ---------------------------------------------------------------------------
|
||
async def handle_user_input(user_input: str) -> str:
|
||
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:
|
||
user_input = input("> ")
|
||
if not user_input:
|
||
continue
|
||
result = await handle_user_input(user_input)
|
||
if result == "quit":
|
||
print("Goodbye!")
|
||
break
|
||
print(result)
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|