query_graph.py

import json
from collections import Counter
import os
import time
import logging
import numpy as np
from openai import OpenAI
from tools.utils import truncate_text
from database_utils import (
    search_vector_search,
    fetch_paths_between_entities,
    search_community,
    get_text_units,
    get_entity_levels,
    get_embedding,
)
from prompt import PROMPTS
from itertools import combinations
from dotenv import load_dotenv

load_dotenv()
logger=logging.getLogger(__name__)

# Load configuration from environment variables
MODEL = os.getenv('MODEL_LLM', 'grok-4-fast-reasoning')
EMBEDDING_MODEL = os.getenv('TOGETHER_MODEL', 'intfloat/multilingual-e5-large-instruct')
EMBEDDING_URL = os.getenv('TOGETHER_EMBEDDING_URL', 'https://api.together.xyz/embedding')

def embedding(texts: list[str]) -> np.ndarray:
    # """doublette, see also embedding in build_graph and get_embedding in database_utils"""
    # model_name = EMBEDDING_MODEL
    # api_key = os.getenv('TOGETHER_API_KEY')
    # client = OpenAI(
    #     api_key=api_key,
    #     base_url=EMBEDDING_URL
    # ) 
    # embedding = client.embeddings.create(
    #     input=[truncate_text(txt, max_tokens=int(os.getenv('TOGETHER_EMBED_MAX_TOKENS', '480')))[ : int(os.getenv('TOGETHER_EMBED_MAX_TOKENS', '480')) * 3] for txt in texts],
    #     model=model_name,
    # )
    # final_embedding = [d.embedding for d in embedding.data]
    # return np.array(final_embedding)
    return np.array(get_embedding(texts))

def generate_llm_response(query, system_prompt):
    """Generate LLM response using OpenAI API"""
    api_key = os.getenv('OPENAI_API_KEY')
    base_url = os.getenv('OPENAI_BASE_URL', 'https://api.openai.com/v1')
    
    if not api_key:
        raise ValueError("OPENAI_API_KEY environment variable not set")
    
    client = OpenAI(api_key=api_key, base_url=base_url)
    
    try:
        completion = client.chat.completions.create(
            model=MODEL,
            messages=[
                {'role': 'system', 'content': system_prompt},
                {'role': 'user', 'content': query}
            ],
            max_tokens=4096,
            temperature=0
        )
        
        return completion.choices[0].message.content
    
    except Exception as e:
        print(f"Error calling OpenAI API: {e}")
        return ""

## Removed duplicate truncate_text and config block (using earlier definitions)

def get_reasoning_chain(global_config, entities_set):
    """Compute shortest paths between candidate entities for reasoning context."""

    maybe_edges = list(combinations(entities_set, 2))
    path_depth = global_config.get('path_depth', 4)
    path_details = fetch_paths_between_entities(maybe_edges, max_depth=path_depth)

    if not path_details:
        return [], ""

    relation_lines = []
    for path in path_details:
        rels = path.get('relationships') or []
        if rels:
            formatted = [
                f"{rel.get('start')} -[{rel.get('description') or 'RELATES_TO'}]-> {rel.get('end')}"
                for rel in rels
            ]
            relation_lines.append("\n".join(formatted))
        else:
            relation_lines.append(" -> ".join(path.get('nodes', [])))

    reasoning_path_information_description = "\n".join(relation_lines)
    return path_details, reasoning_path_information_description

def get_entity_description(global_config, entities, mode=0):
    columns = ['entity_name', 'parent', 'description']
    lines = []
    for item in entities:
        if isinstance(item, dict):
            name = item.get('entity_name', '')
            parent = item.get('parent') or ''
            description = item.get('description', '')
        else:
            name = item[0]
            parent = item[1] or ''
            description = item[2]
        lines.append(f"{name}\t\t{parent}\t\t{description}")

    entity_descriptions = "\t\t".join(columns) + "\n" + "\n".join(lines)
    return entity_descriptions
        
def get_aggregation_description(global_config, reasoning_paths, if_findings=False):
    aggregation_results = []
    visited_nodes = set()

    for path in reasoning_paths:
        visited_nodes.update(path.get('nodes', []))

    for community in visited_nodes:
        temp = search_community(community, None)
        if not temp:
            continue
        aggregation_results.append(temp)

    if if_findings:
        columns = ['entity_name', 'entity_description', 'findings']
        aggregation_descriptions = "\t\t".join(columns) + "\n"
        aggregation_descriptions += "\n".join(
            [information[0] + "\t\t" + str(information[1]) + "\t\t" + information[2] for information in aggregation_results]
        )
    else:
        columns = ['entity_name', 'entity_description']
        aggregation_descriptions = "\t\t".join(columns) + "\n"
        aggregation_descriptions += "\n".join(
            [information[0] + "\t\t" + str(information[1]) for information in aggregation_results]
        )
    
    return aggregation_descriptions, visited_nodes, aggregation_results


def build_reference_summary(vector_hits, path_details, aggregation_records):
    """Create a concise human-readable summary highlighting traversal results."""

    lines: list[str] = []

    if vector_hits:
        lines.append("Matched Entities:")
        for hit in vector_hits:
            name = hit.get('entity_name')
            parent_label = hit.get('parent') or 'root'
            summary_desc = truncate_text(hit.get('description', ''), max_tokens=60)
            snippet = f"- {name} (parent: {parent_label})"
            if summary_desc:
                snippet += f" — {summary_desc}"
            lines.append(snippet)
        lines.append("")

    lines.append("Graph Traversal Paths:")
    if path_details:
        seen_paths = set()
        for path in path_details:
            nodes = path.get('nodes') or []
            path_repr = " -> ".join(nodes)
            if path_repr and path_repr not in seen_paths:
                lines.append(f"- {path_repr}")
                seen_paths.add(path_repr)
            for rel in path.get('relationships') or []:
                desc = rel.get('description') or 'RELATES_TO'
                relation_line = f"    {rel.get('start')} -[{desc}]-> {rel.get('end')}"
                lines.append(relation_line)
        lines.append("")
    else:
        lines.append("- No connecting paths found among the matched entities.\n")

    lines.append("Communities:")
    if aggregation_records:
        for entity_name, description, _ in aggregation_records:
            summary_desc = truncate_text(description or '', max_tokens=60)
            entry = f"- {entity_name}"
            if summary_desc:
                entry += f": {summary_desc}"
            lines.append(entry)
    else:
        lines.append("- No community summaries retrieved.")

    return "\n".join(lines).strip()


def format_layer_counts(counter: Counter) -> str:
    if not counter:
        return "none"
    parts = [f"L{level}: {count}" for level, count in sorted(counter.items())]
    return ", ".join(parts)


def build_layer_overview(vector_hits, path_details, level_lookup):
    entity_level_counts = Counter()
    for hit in vector_hits:
        level = hit.get('level')
        if level is None:
            level = level_lookup.get(hit.get('entity_name'))
        if level is not None:
            entity_level_counts[level] += 1

    path_node_counts = Counter()
    relation_level_counts = Counter()
    for path in path_details:
        for node in path.get('nodes', []):
            level = level_lookup.get(node)
            if level is not None:
                path_node_counts[level] += 1
        for rel in path.get('relationships', []):
            level = rel.get('level')
            if level is not None:
                relation_level_counts[level] += 1

    lines = ["Layer Touch Summary:"]
    lines.append(f"  • Retrieved entities: {format_layer_counts(entity_level_counts)}")
    lines.append(f"  • Path nodes traversed: {format_layer_counts(path_node_counts)}")
    lines.append(f"  • Relation levels: {format_layer_counts(relation_level_counts)}")
    return "\n".join(lines)


def query_graph(global_config,db,query):
    embedding: callable=global_config["embeddings_func"]
    b=time.time()
    level_mode=global_config['level_mode']
    topk=global_config['topk']
    chunks_file=global_config["chunks_file"]
    working_dir=global_config.get("working_dir", None)
    raw_hits = search_vector_search(working_dir, embedding([query]), topk=topk, level_mode=level_mode)
    v=time.time()
    vector_hits = []
    for row in raw_hits:
        name, parent, description, source_id = row[:4]
        level = row[4] if len(row) > 4 else None
        vector_hits.append({
            'entity_name': name,
            'parent': parent,
            'description': description,
            'source_id': source_id,
            'level': level,
        })

    res_entity = [hit['entity_name'] for hit in vector_hits]
    chunks = [hit['source_id'] for hit in vector_hits if hit.get('source_id')]
    entity_descriptions = get_entity_description(global_config, vector_hits)
    try:
        path_details, reasoning_context = get_reasoning_chain(global_config, res_entity)
    except Exception as e:
        print(f"Warning: Reasoning chain failed: {e}")
        path_details = []
        reasoning_context = ""

    try:
        aggregation_descriptions, _, aggregation_records = get_aggregation_description(
            global_config,
            path_details,
        )
    except Exception as e:
        print(f"Warning: Aggregation description failed: {e}")
        aggregation_descriptions = ""
        aggregation_records = []

    text_unit_previews = get_text_units(working_dir, chunks, None, k=5, max_length=240)

    node_names = set(res_entity)
    for path in path_details:
        node_names.update(path.get('nodes', []))
    level_lookup = get_entity_levels(list(node_names))
    layer_overview = build_layer_overview(vector_hits, path_details, level_lookup)

    context_sections = [
        f"entity_information:\n{entity_descriptions}",
        f"aggregation_entity_information:\n{aggregation_descriptions}",
        layer_overview,
    ]
    if reasoning_context:
        context_sections.append(f"reasoning_path_information:\n{reasoning_context}")
    if text_unit_previews:
        context_sections.append("text_units:\n" + "\n".join(text_unit_previews))

    describe = "\n".join(section.strip() for section in context_sections if section)
    e=time.time()
    
    # print(describe)
    sys_prompt =PROMPTS["rag_response"].format(context_data=describe)
    response=generate_llm_response(query, sys_prompt)
    g=time.time()
    print(f"embedding time: {v-b:.2f}s")
    print(f"query time: {e-v:.2f}s")
    
    print(f"response time: {g-e:.2f}s")
    summary = "\n".join([
        layer_overview,
        build_reference_summary(vector_hits, path_details, aggregation_records)
    ])
    return summary, response

# Note: Database connections are now handled through database_utils.py
# using SQLite for relational data and Memgraph for graph operations