from langchain.tools import tool
from qdrant_client import QdrantClient
from langchain.embeddings.ollama import OllamaEmbeddings
from langchain.vectorstores.qdrant import QdrantVectorStore
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import Document

# Initialize embeddings and vector store
embeddings = OllamaEmbeddings(model="nomic-embed-text")
client = QdrantClient(host="localhost", port=6333)
collection_name = "knowledge_base"
# Ensure collection exists
if not client.has_collection(collection_name):
    # Determine embedding dimension by embedding a dummy text
    dim = embeddings.embed_query(["test"])[0].shape[0]
    client.create_collection(name=collection_name, vectors_config={"size": dim, "distance": "Cosine"})
vector_store = QdrantVectorStore(client=client, collection_name=collection_name, embedding=embeddings)

# Text splitter
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)

@tool("search_knowledge_base")
def search_knowledge_base(query: str, max_results: int = 5) -> str:
    """Search the knowledge base for relevant documents."""
    results = vector_store.similarity_search_with_score(query, k=max_results)
    return "\n---\n".join([f"{score:.4f}: {doc.page_content[:200]}..." for doc, score in results])

@tool("add_to_knowledge_base")
def add_to_knowledge_base(content: str, title: str) -> str:
    """Add a new document to the knowledge base."""
    chunks = splitter.split_text(content)
    docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks]
    vector_store.add_documents(docs)
    return f"Added {len(chunks)} chunks from '{title}'."