DRIFT Search

In [1]:

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

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

In [2]:

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

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

---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
Cell In[2], line 19
     15 from graphrag.tokenizer.get_tokenizer import get_tokenizer
     16 from graphrag_llm.completion import create_completion
     17 from graphrag_llm.config import ModelConfig
     18 from graphrag_llm.embedding import create_embedding
---> 19 from graphrag_vectors import IndexSchema, LanceDBVectorStore

ImportError: cannot import name 'LanceDBVectorStore' from 'graphrag_vectors' (/home/runner/work/graphrag/graphrag/packages/graphrag-vectors/graphrag_vectors/__init__.py)

In [3]:

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)

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)

In [4]:

# 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)}"
)

# 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)}" )

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[4], line 23
     19 relationships = read_indexer_relationships(relationship_df)
     20 text_units = read_indexer_text_units(text_unit_df)
     21 
     22 # Connect to the existing GraphRAG vector index for entity-description embeddings.
---> 23 description_embedding_store = LanceDBVectorStore(
     24     index_schema=IndexSchema(index_name="default-entity-description")
     25 )
     26 description_embedding_store.connect(db_uri=LANCEDB_URI)

NameError: name 'LanceDBVectorStore' is not defined

In [5]:

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
)

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 )

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[5], line 13
      9     text_embedder=text_embedder,
     10     entities=entities,
     11     relationships=relationships,
     12     reports=reports,
---> 13     entity_text_embeddings=description_embedding_store,
     14     text_units=text_units,
     15     tokenizer=tokenizer,
     16     config=drift_params,

NameError: name 'description_embedding_store' is not defined

In [6]:

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

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

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[6], line 1
----> 1 resp = await search.search("Who is agent Mercer?")

NameError: name 'search' is not defined

In [7]:

resp.response

resp.response

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[7], line 1
----> 1 resp.response

NameError: name 'resp' is not defined

In [8]:

print(resp.context_data)

print(resp.context_data)

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[8], line 1
----> 1 print(resp.context_data)

NameError: name 'resp' is not defined