Удалить rag_agent.py

This commit is contained in:
2026-05-28 13:26:19 +00:00
parent 3a686e4bb7
commit a42faf40d1
-99
View File
@@ -1,99 +0,0 @@
"""
RAG agent implementation with Qdrant and Ollama.
"""
import os
from pathlib import Path
from typing import List, Dict
from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain.tools import tool
from langchain_core.messages import HumanMessage
from langchain.agents import create_agent
# Initialize LLM and embeddings using Ollama
LLM_MODEL = "llama3"
EMBEDDING_MODEL = "nomic-embed-text"
llm = ChatOllama(model=LLM_MODEL, temperature=0.0)
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
# Qdrant client (inmemory for simplicity)
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams
client = QdrantClient(":memory:")
COLLECTION_NAME = "knowledge"
if not client.collection_exists(COLLECTION_NAME):
client.create_collection(
COLLECTION_NAME,
vectors_config=VectorParams(size=embeddings.embedding_size, distance=Distance.COSINE),
)
vector_store = QdrantVectorStore(client=client, collection_name=COLLECTION_NAME, embedding=embeddings)
# Tool: add to knowledge base
@tool
def add_to_knowledge_base(content: str, title: str = "document") -> str:
"""Add a document to the vector store.
Parameters
----------
content: str
Raw text of the document.
title: str, optional
Title or identifier for the document.
Returns
-------
str
Confirmation message.
"""
# Split into chunks using chunker module
from chunker import split_text
chunks = split_text(content)
docs = []
for i, chunk in enumerate(chunks):
meta = {"title": title, "chunk_index": str(i)}
docs.append({"page_content": chunk, "metadata": meta})
vector_store.add_documents(docs)
return f"Added {len(chunks)} chunks from '{title}'."
# Tool: search knowledge base
@tool
def search_knowledge_base(query: str, max_results: int = 5) -> str:
"""Semantic search in the vector store.
Parameters
----------
query: str
Search query.
max_results: int, optional
Number of top results to return.
Returns
-------
str
Formatted search results.
"""
docs = vector_store.similarity_search(query, k=max_results)
if not docs:
return "No relevant documents found."
lines = []
for i, doc in enumerate(docs, 1):
title = doc.metadata.get("title", "unknown")
chunk_idx = doc.metadata.get("chunk_index", "0")
lines.append(f"{i}. [{title} - chunk {chunk_idx}]\n{doc.page_content[:200]}...")
return "\n\n".join(lines)
# Create agent with tools
SYSTEM_PROMPT = (
"You are an assistant that can search and add documents to a knowledge base."
" Use the provided tools to manage the knowledge base."
)
agent = create_agent(
llm=llm,
tools=[add_to_knowledge_base, search_knowledge_base],
system_prompt=SYSTEM_PROMPT,
)
# Expose agent for external use
__all__ = ["agent", "add_to_knowledge_base", "search_knowledge_base"]