Copyright (c) 2024 Microsoft Corporation.

# Copyright (c) 2024 Microsoft Corporation.
# Licensed under the MIT License.

import os

import pandas as pd
from graphrag.config.models.drift_search_config import DRIFTSearchConfig
from graphrag.query.indexer_adapters import (
    read_indexer_entities,
    read_indexer_relationships,
    read_indexer_reports,
    read_indexer_text_units,
)
from graphrag.query.structured_search.drift_search.drift_context import (
    DRIFTSearchContextBuilder,
)
from graphrag.query.structured_search.drift_search.search import DRIFTSearch
from graphrag.tokenizer.get_tokenizer import get_tokenizer
from graphrag_llm.completion import create_completion
from graphrag_llm.config import ModelConfig
from graphrag_llm.embedding import create_embedding
from graphrag_vectors import IndexSchema, LanceDBVectorStore

api_key = os.environ["GRAPHRAG_API_KEY"]

chat_config = ModelConfig(
    type="litellm",
    model_provider="openai",
    model="gpt-4.1",
    api_key=api_key,
)
chat_model = create_completion(chat_config)

tokenizer = get_tokenizer(chat_config)

embedding_config = ModelConfig(
    type="litellm",
    model_provider="openai",
    model="text-embedding-3-large",
    api_key=api_key,
)

text_embedder = create_embedding(embedding_config)

# parquet files generated from indexing pipeline
INPUT_DIR = "./inputs/operation dulce"
LANCEDB_URI = "./lancedb"
COMMUNITY_TABLE = "communities"
COMMUNITY_REPORT_TABLE = "community_reports"
ENTITY_TABLE = "entities"
RELATIONSHIP_TABLE = "relationships"
TEXT_UNIT_TABLE = "text_units"
COMMUNITY_LEVEL = 2

community_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_TABLE}.parquet")
report_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet")
entity_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_TABLE}.parquet")
relationship_df = pd.read_parquet(f"{INPUT_DIR}/{RELATIONSHIP_TABLE}.parquet")
text_unit_df = pd.read_parquet(f"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet")

reports = read_indexer_reports(report_df, community_df, COMMUNITY_LEVEL)
entities = read_indexer_entities(entity_df, community_df, COMMUNITY_LEVEL)
relationships = read_indexer_relationships(relationship_df)
text_units = read_indexer_text_units(text_unit_df)

# Connect to the existing GraphRAG vector index for entity-description embeddings.
description_embedding_store = LanceDBVectorStore(
    index_schema=IndexSchema(index_name="default-entity-description")
)
description_embedding_store.connect(db_uri=LANCEDB_URI)

print(
    f"Loaded reports={len(reports)}, entities={len(entities)}, relationships={len(relationships)}, text_units={len(text_units)}"
)

drift_params = DRIFTSearchConfig(
    primer_folds=1,
    drift_k_followups=3,
    n_depth=3,
)

context_builder = DRIFTSearchContextBuilder(
    model=chat_model,
    text_embedder=text_embedder,
    entities=entities,
    relationships=relationships,
    reports=reports,
    entity_text_embeddings=description_embedding_store,
    text_units=text_units,
    tokenizer=tokenizer,
    config=drift_params,
)

search = DRIFTSearch(
    model=chat_model, context_builder=context_builder, tokenizer=tokenizer
)

resp = await search.search("Who is agent Mercer?")

resp.response

print(resp.context_data)