vectorstore.py

from typing import List, Dict, Any, Optional, Tuple
import numpy as np
from qdrant_client import QdrantClient
from qdrant_client.http import models

import sys, os
sys.path.append(os.path.dirname( os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
from config.envconfig import envconfig
import json
import os
from agent_workflow.state import Triplet
import logging 

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class QdrantConnector:
   
    
    def __init__(self, 
                qdrant_url: Optional[str] = None, 
                api_key: Optional[str] = None,
                local_path: str = "./qdrant_data",
                vector_size: int = 1536):  
      
        if qdrant_url and api_key:
            self.client = QdrantClient(
                url=qdrant_url, 
                api_key=api_key,
                timeout = 50, 
                )
            
            
            self.using_cloud = True
        else:
            self.client = QdrantClient(path=local_path)
            self.using_cloud = False
            
        self.vector_size = vector_size
        self.collection_name = "unique_triplets"
        self.node_collection_name = "unique_nodes"
        self.baseline_collection_name = "baseline_timenet"
    
    def initialize_collections(self):
        
        try:
            
            collections = self.client.get_collections()
            collection_names = [c.name for c in collections.collections]

            
            for name in [self.collection_name, self.node_collection_name, self.baseline_collection_name]:
                if name not in collection_names:
                    self.client.create_collection(
                        collection_name=name,
                        vectors_config=models.VectorParams(
                            size=self.vector_size,
                            distance=models.Distance.COSINE
                        )
                    )
                    logger.info(f"Created Qdrant collection: {name}")
                else:
                    logger.info(f"Collection {name} does not exist, will create it.")

            logger.info("Update collection indexing ...")
            self.update_collection()

        except Exception as e:
            logger.error(f"Error initializing Qdrant collection: {e}")

    def store_triplets(self, triplet_embeddings: List[Tuple[Dict[str, Any], np.ndarray]]) -> int:
     
        if not triplet_embeddings:
            return 0
        
        points = []
        
        for i, (triplet, embedding) in enumerate(triplet_embeddings):
            point_id = f"{triplet['subject_id']}_{triplet['predicate']}_{triplet['object_id']}"
            point_id_hash = abs(hash(point_id)) % (2**63 - 1)  # Convert to positive integer
            
            payload = {
                "subject_id": triplet["subject_id"],
                "subject_name": triplet["subject_name"],
                "subject_type": triplet.get("subject_type", "Entity"),
                "predicate": triplet["predicate"],
                "object_id": triplet["object_id"],
                "object_name": triplet["object_name"],
                "object_type": triplet.get("object_type", "Entity"),
                "description": triplet.get("description", "")
            }
            
            points.append(
                models.PointStruct(
                    id=point_id_hash,
                    vector=embedding,
                    payload=payload
                )
            )

        try:
            batch_size = 100
            for i in range(0, len(points), batch_size):
                batch = points[i:i+batch_size]
                self.client.upsert(
                    collection_name=self.collection_name,
                    points=batch
                )
            return len(points)
        except Exception as e:
            logger.error(f"Error storing triplets in Qdrant: {e}")
            return 0
    
    
    def search_triplets(self, query_vector: np.ndarray, limit: int = 10) -> List[Dict[str, Any]]:
       
        try:
            results = self.client.search(
                collection_name=self.collection_name,
                query_vector=query_vector,
                limit=limit
            )
            
            formatted_results = []
            for hit in results:
                formatted_results.append({
                    "triplet": f"{hit.payload['subject_name']} {hit.payload['predicate']} {hit.payload['object_name']}",
                    "subject_id": hit.payload["subject_id"],
                    "subject_name": hit.payload["subject_name"],
                    "predicate": hit.payload["predicate"],
                    "object_id": hit.payload["object_id"],
                    "object_name": hit.payload["object_name"],
                    "description": hit.payload.get("description", ""),
                    "score": hit.score
                })
            
            return formatted_results
        except Exception as e:
            logger.error(f"Error searching Qdrant: {e}")
            return []
    
    def delete_collections(self):
        """
        Delete all collections in Qdrant.
        """

        try:
            self.client.recreate_collection(
                collection_name=self.collection_name,
                vectors_config=models.VectorParams(
                    size=self.vector_size,
                    distance=models.Distance.COSINE
                )
            )
            logger.info(f"Deleted Qdrant collection: {self.collection_name}")
        except Exception as e:
            logger.error(f"Error deleting Qdrant collection: {e}")


    def store_nodes(self, nodes_embeddings: List[Tuple[Dict[str, Any], np.ndarray]]) -> int:
        
        if not nodes_embeddings:
            return 0
        
        points = []
        
        for i, (node, embedding) in enumerate(nodes_embeddings):
            point_id = node["id"]
            point_id_hash = abs(hash(point_id)) % (2**63 - 1)
            payload = {
                "id": node["id"],
                "name": node["name"],
                "type": node.get("type", "Entity"),
                "description": node.get("description", "")
            }
            points.append(
                models.PointStruct(
                    id=point_id_hash,
                    vector=embedding,
                    payload=payload
                )
            )
        try:
            batch_size = 100
            for i in range(0, len(points), batch_size):
                batch = points[i:i+batch_size]
                self.client.upsert(
                    collection_name=self.node_collection_name,
                    points=batch
                )
            return len(points)
        except Exception as e:
            logger.error(f"Error storing nodes in Qdrant: {e}")
            return 0


    def search_nodes(self, query_vector: np.ndarray, limit: int = 10, threshold: float = 0.5) -> List[Dict[str, Any]]:
        try:
            results = self.client.search(
                collection_name=self.node_collection_name,
                query_vector=query_vector,
                limit=limit, 
                score_threshold=threshold
            )
            
            formatted_results = []
            for hit in results:
                formatted_results.append({
                    "id": hit.payload["id"],
                    "name": hit.payload["name"],
                    "type": hit.payload.get("type", "Entity"),
                    "description": hit.payload.get("description", ""),
                    "score": hit.score
                })
            
            return formatted_results
        except Exception as e:
            print(f"Error searching Qdrant nodes: {e}")
            return []
    

    def search_triplets_by_query(self, embedding: np.ndarray, triplets : List[Triplet], limit: int = 10) -> List[Dict[str, Any]]:
        """
        Search for triplets in Qdrant by their embeddings.
        """
        
        
        filter = models.Filter(
            must=[],
            should=[],
            must_not=[]
        )
        for triplet in triplets:
            triplet_condition = models.Filter(
                    must=[
                        models.FieldCondition(
                            key="subject_id",
                            match=models.MatchValue(value=triplet.subject_id)
                        ),
                        models.FieldCondition(
                            key="predicate",
                            match=models.MatchValue(value=triplet.predicate)
                        ),
                        models.FieldCondition(
                            key="object_id",
                            match=models.MatchValue(value=triplet.object_id)
                        )
                    ]
                )
            
            filter.should.append(triplet_condition)

        all_results = self.client.search(
            collection_name=self.collection_name,
            query_vector=embedding,
            limit=limit, 
            # query_filter=filter, 
            with_payload=True
        )

        return all_results



    def update_collection(self): 
        """
        Update the collection with necessary payload indexes.
        """

        self.client.create_payload_index(
            collection_name=self.collection_name,
            field_name="subject_id",
            field_schema="keyword"
        )

        self.client.create_payload_index(
            collection_name=self.collection_name,
            field_name="predicate",
            field_schema="keyword"
        )

        self.client.create_payload_index(
            collection_name=self.collection_name,
            field_name="object_id",
            field_schema="keyword"
        )

        


    def store_event(self, event_embeddings: List[Tuple[str, np.ndarray]]) -> int: 

        points = []
        for i, (event_id, embedding) in enumerate(event_embeddings):
            point_id = f"event_{i}"
            point_id_hash = abs(hash(point_id)) % (2**63 - 1)
            payload = {
                "event": event_id
            }
            points.append(
                models.PointStruct(
                    id=point_id_hash,
                    vector=embedding,  
                    payload=payload
                )
            )

        try:
            batch_size = 100
            for i in range(0, len(points), batch_size):
                batch = points[i:i+batch_size]
                self.client.upsert(
                    collection_name=self.baseline_collection_name,
                    points=batch
                )
            return len(points)
        except Exception as e:
            print(f"Error storing events in Qdrant: {e}")
            return 0


    def search_events(self, query_vector: np.ndarray, limit: int = 10) -> str:
        
        try:
            results = self.client.search(
                collection_name=self.baseline_collection_name,
                query_vector=query_vector,
                limit=limit
            )
            
            formatted_results = ""
            for i, hit in enumerate(results):
                formatted_results += f"\n Chunk{i} Event: {hit.payload['event']}"

            return formatted_results
        except Exception as e:
            print(f"Error searching Qdrant events: {e}")
            return []



if __name__ == "__main__":

    from data_processing.data_transform.embedder import TripletEmbedder


    # Example usage
    connector = QdrantConnector(
        qdrant_url=envconfig.QDRANT_URL, 
        api_key=envconfig.QDRANT_API_KEY
    )

    connector.initialize_collections()

    query = "Sạt lở đất và lũ quét năm 2024"
    embedder = TripletEmbedder()
    query_embedding = embedder.embed_text(query)
    triplet_string = [
        {"subject_id":"landslide_flood_2024","subject_name":"Sạt lở đất và lũ quét năm 2024","subject_type":"NEWS_EVENT","predicate":"CAUSED","object_id":"yen_bai_damage_2024","object_name":"Thiệt hại tại Yên Bái năm 2024","object_type":"NEWS_EVENT","description":"Sạt lở đất và lũ quét năm 2024 gây thiệt hại lớn tại Yên Bái"},
        {"subject_id":"viettel_double_digit_growth_7_years","subject_name":"Viettel đạt tăng trưởng 2 con số sau 7 năm","subject_type":"NEWS_EVENT","predicate":"PRECEDES","object_id":"viettel_growth_2024","object_name":"Viettel tăng trưởng 2 con số năm 2024","object_type":"NEWS_EVENT","description":"Lần đầu tiên sau 7 năm, Viettel đạt tăng trưởng 2 con số trong năm 2024."}, 
    ]

    def parse_triplets_from_json(triplet_json_strings: List[Dict[str, Any]]) -> List[Triplet]:
        triplets = []

        for triplet_dict in triplet_json_strings:
            try:
               
                triplet = Triplet(
                    subject_id=triplet_dict.get("subject_id", ""),
                    subject_name=triplet_dict.get("subject_name", ""),
                    subject_type=triplet_dict.get("subject_type", "Entity"),
                    predicate=triplet_dict.get("predicate", ""),
                    object_id=triplet_dict.get("object_id", ""),
                    object_name=triplet_dict.get("object_name", ""),
                    object_type=triplet_dict.get("object_type", "Entity") if triplet_dict.get("object_type") else "Entity",
                    description=triplet_dict.get("description", "")
                )
                triplets.append(triplet)
            except json.JSONDecodeError as e:
                print(f"Error parsing triplet JSON: {e}")
            except Exception as e:
                print(f"Error creating Triplet object: {e}")
        
        return triplets


    triplets = parse_triplets_from_json(triplet_string)

    result = connector.search_triplets_by_query(
        embedding=query_embedding, 
        triplets=triplets, 
        limit=10
    )
    for hit in result:
        print(f"Triplet: {hit.payload['subject_name']} {hit.payload['predicate']} {hit.payload['object_name']}, Score: {hit.score}")

    print("Done!")