import logging from typing import List, Dict, Any from qdrant_client import QdrantClient, models from qdrant_client.models import VectorParams, Distance, PointStruct, Filter, FieldCondition, Range from sentence_transformers import SentenceTransformer from fastapi import FastAPI, HTTPException from pydantic import BaseModel # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') QDRANT_HOST = "qdrant-b3knu-u50607.vm.elestio.app" QDRANT_PORT = 6333 COLLECTION_NAME = "genesis_knowledge" EMBEDDING_MODEL = "all-mpnet-base-v2" # A good general-purpose model # Define data models for API requests and responses class SearchRequest(BaseModel): query: str limit: int = 10 filters: Dict[str, Any] = {} # Example: {"entity_type": "person", "confidence": {"gte": 0.8}} class SearchResult(BaseModel): id: str score: float payload: Dict[str, Any] class KnowledgeSearchIndex: """ Builds and manages a search index over extracted knowledge using Qdrant. Supports full-text search, entity search, and concept search via vector similarity. """ def __init__(self, qdrant_host: str = QDRANT_HOST, qdrant_port: int = QDRANT_PORT, collection_name: str = COLLECTION_NAME, embedding_model: str = EMBEDDING_MODEL): """ Initializes the KnowledgeSearchIndex. Args: qdrant_host: The hostname of the Qdrant instance. qdrant_port: The port of the Qdrant instance. collection_name: The name of the Qdrant collection to use. embedding_model: The name of the SentenceTransformer model to use for embeddings. """ self.qdrant_host = qdrant_host self.qdrant_port = qdrant_port self.collection_name = collection_name self.embedding_model = embedding_model self.client = QdrantClient(host=self.qdrant_host, port=self.qdrant_port) self.model = SentenceTransformer(self.embedding_model) try: self.client.get_collection(collection_name=self.collection_name) logging.info(f"Collection '{self.collection_name}' already exists.") except Exception as e: logging.info(f"Collection '{self.collection_name}' does not exist, creating it.") self.create_collection() def create_collection(self): """ Creates the Qdrant collection if it doesn't exist. """ try: self.client.recreate_collection( collection_name=self.collection_name, vectors_config=VectorParams(size=self.model.get_sentence_embedding_dimension(), distance=Distance.COSINE), ) logging.info(f"Collection '{self.collection_name}' created successfully.") except Exception as e: logging.error(f"Error creating collection '{self.collection_name}': {e}") raise def index_knowledge(self, knowledge_items: List[Dict[str, Any]]): """ Indexes a list of knowledge items in Qdrant. Args: knowledge_items: A list of dictionaries, where each dictionary represents a knowledge item. Each dictionary should have a 'text' key containing the text to be indexed, and any other keys will be stored as payload. A unique 'id' field is required. """ points = [] for item in knowledge_items: try: text = item.get("text") item_id = item.get("id") if not text: logging.warning(f"Skipping knowledge item with missing 'text': {item}") continue if not item_id: logging.warning(f"Skipping knowledge item with missing 'id': {item}") continue embedding = self.model.encode(text).tolist() payload = {k: v for k, v in item.items() if k != "text"} # Store all other fields as payload point = PointStruct(id=item_id, vector=embedding, payload=payload) points.append(point) except Exception as e: logging.error(f"Error indexing knowledge item: {item}. Error: {e}") try: self.client.upsert(collection_name=self.collection_name, points=points, wait=True) logging.info(f"Indexed {len(points)} knowledge items.") except Exception as e: logging.error(f"Error upserting points to Qdrant: {e}") raise def search(self, query: str, limit: int = 10, filters: Dict[str, Any] = {}) -> List[SearchResult]: """ Searches the knowledge base for items similar to the query. Args: query: The search query. limit: The maximum number of results to return. filters: A dictionary of filters to apply to the search. Example: {"entity_type": "person", "confidence": {"gte": 0.8}} Supports equality and range filters. Returns: A list of SearchResult objects, sorted by relevance. """ try: query_vector = self.model.encode(query).tolist() qdrant_filter = self._build_qdrant_filter(filters) search_result = self.client.search( collection_name=self.collection_name, query_vector=query_vector, limit=limit, query_filter=qdrant_filter ) results = [ SearchResult(id=hit.id, score=hit.score, payload=hit.payload) for hit in search_result ] return results except Exception as e: logging.error(f"Error during search: {e}") raise def _build_qdrant_filter(self, filters: Dict[str, Any]) -> models.Filter: """ Builds a Qdrant filter from a dictionary of filters. Args: filters: A dictionary of filters. Example: {"entity_type": "person", "confidence": {"gte": 0.8}} Returns: A Qdrant Filter object. """ filter_list = [] for key, value in filters.items(): if isinstance(value, dict): # Range filter if "gte" in value: filter_list.append(FieldCondition(key=key, range=Range(gte=value["gte"]))) if "lte" in value: filter_list.append(FieldCondition(key=key, range=Range(lte=value["lte"]))) else: # Equality filter filter_list.append(FieldCondition(key=key, match=models.MatchValue(value=value))) if filter_list: return Filter(must=filter_list) else: return None # FastAPI integration app = FastAPI() knowledge_index = KnowledgeSearchIndex() @app.post("/api/knowledge/search", response_model=List[SearchResult]) async def search_knowledge(search_request: SearchRequest): """ Searches the knowledge base. """ try: results = knowledge_index.search(search_request.query, search_request.limit, search_request.filters) return results except Exception as e: logging.error(f"Error during API search: {e}") raise HTTPException(status_code=500, detail=str(e)) if __name__ == "__main__": # Example usage: # 1. Start Qdrant: docker run -d -p 6333:6333 -p 6334:6334 qdrant/qdrant # 2. Run this script. # 3. Send a POST request to /api/knowledge/search with a JSON body like: # {"query": "What is the capital of France?", "limit": 5} # or # {"query": "Find people with high confidence", "limit": 5, "filters": {"entity_type": "person", "confidence": {"gte": 0.8}}} # Example knowledge items knowledge_items = [ {"id": "1", "text": "The capital of France is Paris.", "entity_type": "location"}, {"id": "2", "text": "Albert Einstein was a famous physicist.", "entity_type": "person", "confidence": 0.95}, {"id": "3", "text": "Quantum mechanics is a fundamental theory in physics.", "entity_type": "concept"}, {"id": "4", "text": "Paris is a beautiful city.", "entity_type": "location"}, {"id": "5", "text": "Marie Curie discovered radium.", "entity_type": "person", "confidence": 0.85}, ] try: knowledge_index.index_knowledge(knowledge_items) # Example search results = knowledge_index.search(query="capital of France", limit=2) print("Search results:", results) # Example search with filters results_filtered = knowledge_index.search(query="famous person", limit=2, filters={"entity_type": "person", "confidence": {"gte": 0.9}}) print("Filtered search results:", results_filtered) except Exception as e: print(f"Error during example usage: {e}")