@susu-eng/gralkor 27.0.0 → 27.1.0

package/server/main.py

@@ -7,6 +7,7 @@ import logging
 import os
 import time
 from contextlib import asynccontextmanager
+from copy import deepcopy
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from typing import Any, Literal
@@ -23,6 +24,7 @@ from graphiti_core.driver.falkordb_driver import FalkorDriver
 from graphiti_core.edges import EntityEdge
 from graphiti_core.nodes import EpisodicNode, EpisodeType
 from graphiti_core.llm_client import LLMConfig
+from graphiti_core.search.search_config_recipes import COMBINED_HYBRID_SEARCH_CROSS_ENCODER
 
 
 # ── Config ────────────────────────────────────────────────────
@@ -134,16 +136,14 @@ def _build_ontology(
     dict[tuple[str, str], list[str]] | None,
     list[str] | None,
 ]:
-    """Build ontology from config. Returns (entity_types, edge_types, edge_type_map, excluded)."""
+    """Build ontology from config. Returns (entity_types, edge_types, edge_type_map)."""
     raw = cfg.get("ontology")
     if not raw:
-        return None, None, None, None
+        return None, None, None
 
     entity_defs = raw.get("entities") or {}
     edge_defs = raw.get("edges") or {}
     edge_map_raw = raw.get("edgeMap") or {}
-    excluded_raw = raw.get("excludedEntityTypes")
-
     entity_types = _build_type_defs(entity_defs) if entity_defs else None
     edge_types = _build_type_defs(edge_defs) if edge_defs else None
 
@@ -154,12 +154,10 @@ def _build_ontology(
             parts = key.split(",")
             edge_type_map[(parts[0], parts[1])] = values
 
-    excluded = list(excluded_raw) if excluded_raw else None
-
-    if not entity_types and not edge_types and not edge_type_map and not excluded:
-        return None, None, None, None
+    if not entity_types and not edge_types and not edge_type_map:
+        return None, None, None
 
-    return entity_types, edge_types, edge_type_map, excluded
+    return entity_types, edge_types, edge_type_map
 
 
 def _log_falkordblite_diagnostics(error: Exception) -> None:
@@ -200,12 +198,11 @@ graphiti: Graphiti | None = None
 ontology_entity_types: dict[str, type[BaseModel]] | None = None
 ontology_edge_types: dict[str, type[BaseModel]] | None = None
 ontology_edge_type_map: dict[tuple[str, str], list[str]] | None = None
-ontology_excluded: list[str] | None = None
 
 
 @asynccontextmanager
 async def lifespan(_app: FastAPI):
-    global graphiti, ontology_entity_types, ontology_edge_types, ontology_edge_type_map, ontology_excluded
+    global graphiti, ontology_entity_types, ontology_edge_types, ontology_edge_type_map
     cfg = _load_config()
 
     falkordb_uri = os.getenv("FALKORDB_URI")
@@ -260,7 +257,7 @@ async def lifespan(_app: FastAPI):
         handler.setFormatter(logging.Formatter("%(message)s"))
         logger.addHandler(handler)
 
-    ontology_entity_types, ontology_edge_types, ontology_edge_type_map, ontology_excluded = _build_ontology(cfg)
+    ontology_entity_types, ontology_edge_types, ontology_edge_type_map = _build_ontology(cfg)
     if ontology_entity_types or ontology_edge_types:
         entity_names = list(ontology_entity_types or {})
         edge_names = list(ontology_edge_types or {})
@@ -282,10 +279,10 @@ def _find_rate_limit_error(exc: Exception) -> Exception | None:
     seen: set[int] = set()
     while current is not None and id(current) not in seen:
         seen.add(id(current))
-        # Match openai.RateLimitError, anthropic.RateLimitError, etc.
-        if type(current).__name__ == "RateLimitError" or (
-            hasattr(current, "status_code") and getattr(current, "status_code", None) == 429
-        ):
+        # Match openai.RateLimitError, anthropic.RateLimitError, google.genai.errors.ClientError, etc.
+        # Note: Google's APIError uses .code, most others use .status_code.
+        http_code = getattr(current, "status_code", None) or getattr(current, "code", None)
+        if type(current).__name__ == "RateLimitError" or http_code == 429:
             return current
         current = current.__cause__ or current.__context__
     return None
@@ -386,6 +383,7 @@ class SearchRequest(BaseModel):
     query: str
     group_ids: list[str]
     num_results: int = 10
+    mode: Literal["fast", "slow"] = "fast"
 
 
 class GroupIdRequest(BaseModel):
@@ -399,6 +397,15 @@ def _ts(dt: datetime | None) -> str | None:
     return dt.isoformat() if dt else None
 
 
+def _serialize_node(node) -> dict[str, Any]:
+    return {
+        "uuid": node.uuid,
+        "name": node.name,
+        "summary": node.summary,
+        "group_id": node.group_id,
+    }
+
+
 def _serialize_fact(edge: EntityEdge) -> dict[str, Any]:
     return {
         "uuid": edge.uuid,
@@ -592,7 +599,7 @@ async def add_episode(req: AddEpisodeRequest):
         entity_types=ontology_entity_types,
         edge_types=ontology_edge_types,
         edge_type_map=ontology_edge_type_map,
-        excluded_entity_types=ontology_excluded,
+        excluded_entity_types=None,
     )
     duration_ms = (time.monotonic() - t0) * 1000
     episode = result.episode
@@ -634,7 +641,7 @@ async def ingest_messages(req: IngestMessagesRequest):
         entity_types=ontology_entity_types,
         edge_types=ontology_edge_types,
         edge_type_map=ontology_edge_type_map,
-        excluded_entity_types=ontology_excluded,
+        excluded_entity_types=None,
     )
     duration_ms = (time.monotonic() - t0) * 1000
     episode = result.episode
@@ -673,33 +680,10 @@ def _ensure_driver_graph(group_ids: list[str] | None) -> None:
         print(f"[gralkor] driver graph routed: {target}", flush=True)
 
 
-def _prioritize_facts(
-    edges: list[EntityEdge], limit: int, reserved_ratio: float = 0.7,
-) -> list[EntityEdge]:
-    """Reserve slots for valid facts, fill the rest by relevance.
-
-    First ~70% of slots are reserved for valid facts (no invalid_at).
-    Remaining slots are filled from whatever Graphiti ranked highest
-    among the leftovers — valid or not — preserving relevance scoring.
-    """
-    reserved_count = max(1, round(limit * reserved_ratio))
-
-    reserved: list[EntityEdge] = []
-    rest: list[EntityEdge] = []
-    for e in edges:
-        if len(reserved) < reserved_count and e.invalid_at is None:
-            reserved.append(e)
-        else:
-            rest.append(e)
-
-    remainder_count = limit - len(reserved)
-    return reserved + rest[:remainder_count]
-
-
 @app.post("/search")
 async def search(req: SearchRequest):
-    logger.info("[gralkor] search — query:%d chars group_ids:%s num_results:%d",
-                len(req.query), req.group_ids, req.num_results)
+    logger.info("[gralkor] search — mode:%s query:%d chars group_ids:%s num_results:%d",
+                req.mode, len(req.query), req.group_ids, req.num_results)
     # graphiti.add_episode() clones the driver to target the correct FalkorDB
     # named graph (database=group_id), but graphiti.search() does not — it just
     # uses whatever graph the driver currently points at. Before the first
@@ -707,26 +691,38 @@ async def search(req: SearchRequest):
     # searches return 0 results. Fix: route to the correct graph here.
     _ensure_driver_graph(req.group_ids)
     t0 = time.monotonic()
-    # Over-fetch to compensate for expired facts that will be deprioritized.
-    fetch_limit = req.num_results * 2
     try:
-        edges = await graphiti.search(
-            query=_sanitize_query(req.query),
-            group_ids=req.group_ids,
-            num_results=fetch_limit,
-        )
+        if req.mode == "slow":
+            # Cross-encoder + BFS: higher quality, also returns entity node summaries.
+            # deepcopy required — COMBINED_HYBRID_SEARCH_CROSS_ENCODER is a module-level
+            # constant; mutating .limit directly would corrupt it across requests.
+            config = deepcopy(COMBINED_HYBRID_SEARCH_CROSS_ENCODER)
+            config.limit = req.num_results
+            search_result = await graphiti.search_(
+                query=_sanitize_query(req.query),
+                group_ids=req.group_ids,
+                config=config,
+            )
+            edges = search_result.edges
+            nodes = search_result.nodes
+        else:
+            edges = await graphiti.search(
+                query=_sanitize_query(req.query),
+                group_ids=req.group_ids,
+                num_results=req.num_results,
+            )
+            nodes = []
     except Exception as e:
         duration_ms = (time.monotonic() - t0) * 1000
-        logger.error("[gralkor] search failed — %.0fms: %s", duration_ms, e)
+        logger.error("[gralkor] search failed — mode:%s %.0fms: %s", req.mode, duration_ms, e)
         raise
     duration_ms = (time.monotonic() - t0) * 1000
-    prioritized = _prioritize_facts(edges, req.num_results)
-    valid_count = sum(1 for e in prioritized if e.invalid_at is None)
-    result = [_serialize_fact(e) for e in prioritized]
-    logger.info("[gralkor] search result — %d facts (%d valid, %d non-valid) from %d fetched %.0fms",
-                len(prioritized), valid_count, len(prioritized) - valid_count, len(edges), duration_ms)
+    result = [_serialize_fact(e) for e in edges]
+    serialized_nodes = [_serialize_node(n) for n in nodes]
+    logger.info("[gralkor] search result — mode:%s %d facts %d nodes %.0fms",
+                req.mode, len(result), len(serialized_nodes), duration_ms)
     logger.debug("[gralkor] search facts: %s", result)
-    return {"facts": result}
+    return {"facts": result, "nodes": serialized_nodes}