@reconcrap/people-network-memory 0.1.0

package/src/people_network_memory/graphiti_adapter/graphiti_store.py

@@ -0,0 +1,655 @@
+"""Graphiti-backed adapter using embedded Kuzu when configured.
+
+Graphiti remains isolated here. The adapter also keeps a JSON projection cache
+for stable person-card output while Graphiti owns semantic episode ingestion and
+graph search.
+"""
+
+from __future__ import annotations
+
+import ast
+import asyncio
+import inspect
+import importlib.util
+import json
+import logging
+import re
+import threading
+import time
+from concurrent.futures import TimeoutError as FutureTimeoutError
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from people_network_memory.config import PeopleMemoryConfig
+from people_network_memory.domain.models import (
+    Evidence,
+    IdentityCandidate,
+    PersonMemoryRecord,
+    PersonRef,
+    RecordInteractionResult,
+    RetrievalItem,
+    SensitivityLabel,
+    SocialInteraction,
+)
+from people_network_memory.graphiti_adapter.episode_formatter import format_graphiti_episode
+from people_network_memory.graphiti_adapter.ontology import EDGE_TYPES, ENTITY_TYPES
+from people_network_memory.infrastructure.embeddings import (
+    EmbeddingSettings,
+    OpenAICompatibleEmbeddingClient,
+)
+from people_network_memory.infrastructure.file_store import JsonPeopleStore
+from people_network_memory.infrastructure.retrieval_intent import (
+    is_follow_up_query,
+    mentioned_query_target,
+    text_answers_mentioned_query,
+)
+from people_network_memory.infrastructure.semantic_index import SemanticProjectionIndex
+from people_network_memory.ports.errors import BackendUnavailableError
+from people_network_memory.ports.errors import SearchError
+
+
+CUSTOM_EXTRACTION_INSTRUCTIONS = """
+Extract personal social-memory entities and relationships conservatively.
+Keep participants separate from people merely mentioned in the note.
+Represent secondhand statements as attributed claims, not direct facts.
+Preserve places, topics, work/school facts, preferences, and follow-up context
+only when supported by the episode body.
+"""
+
+
+class GraphitiGraphStore(JsonPeopleStore):
+    def __init__(self, graphiti: Any, config: PeopleMemoryConfig) -> None:
+        super().__init__(Path(config.data_path).expanduser() / "people-memory.graphiti-cache.json")
+        self._graphiti = graphiti
+        self._config = config
+        self._runner = _AsyncLoopRunner()
+
+    @classmethod
+    def from_config(cls, config: PeopleMemoryConfig) -> "GraphitiGraphStore":
+        if importlib.util.find_spec("graphiti_core") is None:
+            raise BackendUnavailableError(
+                "graphiti_core is not installed. Install with `pip install -e .[graphiti]` "
+                "after the Graphiti spike requirements are confirmed."
+            )
+        graphiti = _build_graphiti(config)
+        return cls(graphiti=graphiti, config=config)
+
+    def find_identity_candidates(self, ref: PersonRef) -> list[IdentityCandidate]:
+        return super().find_identity_candidates(ref)
+
+    def save_interaction(
+        self, interaction: SocialInteraction, identity_map: dict[str, str | None]
+    ) -> RecordInteractionResult:
+        result = super().save_interaction(interaction, identity_map)
+        graphiti_evidence = self.index_interaction_episode(interaction)
+        return result.model_copy(update={"evidence": [*result.evidence, graphiti_evidence]})
+
+    def get_person_memory(self, person_id: str) -> PersonMemoryRecord | None:
+        return super().get_person_memory(person_id)
+
+    def index_interaction_episode(self, interaction: SocialInteraction) -> Evidence:
+        try:
+            graphiti_result = _with_retries(
+                lambda: self._runner.run(
+                    self._add_episode(interaction),
+                    timeout_seconds=self._config.graphiti_add_timeout_seconds,
+                ),
+                label="Graphiti add_episode",
+                attempts=self._config.graphiti_retry_attempts,
+            )
+        except Exception as exc:  # pragma: no cover - live backend defensive path
+            raise BackendUnavailableError(f"Graphiti add_episode failed: {exc}") from exc
+        return Evidence(
+            evidence_id=getattr(graphiti_result.episode, "uuid", "graphiti_episode"),
+            source_text=interaction.source_text,
+            recorded_at=interaction.occurred_at or datetime.now(timezone.utc),
+            confidence=1.0,
+        )
+
+    def cache_interaction_projection(self, interaction: SocialInteraction) -> None:
+        if any(
+            existing.source_text == interaction.source_text
+            and existing.occurred_at == interaction.occurred_at
+            for existing in self.interactions.values()
+        ):
+            return
+        super().save_interaction(interaction, _identity_map_for_cache(self, interaction))
+
+    def close(self) -> None:
+        try:
+            close_graphiti = getattr(self._graphiti, "close", None)
+            if callable(close_graphiti):
+                result = close_graphiti()
+                if inspect.isawaitable(result):
+                    self._runner.run(result, timeout_seconds=10)
+        except TimeoutError:
+            pass
+        finally:
+            self._runner.close()
+
+    def search(
+        self,
+        query: str,
+        *,
+        limit: int = 10,
+        include_sensitive: bool = False,
+        mode: str = "recall",
+    ) -> list[RetrievalItem]:
+        try:
+            edges = _with_retries(
+                lambda: self._runner.run(
+                    self._graphiti.search(query, num_results=limit),
+                    timeout_seconds=self._config.graphiti_search_timeout_seconds,
+                ),
+                label="Graphiti search",
+                attempts=self._config.graphiti_retry_attempts,
+            )
+        except Exception as exc:  # pragma: no cover - live backend defensive path
+            raise SearchError(f"Graphiti search failed: {exc}") from exc
+        items: list[RetrievalItem] = []
+        for edge in edges:
+            fact = getattr(edge, "fact", "")
+            if not include_sensitive and _looks_sensitive(fact):
+                continue
+            evidence = [
+                Evidence(
+                    evidence_id=episode_id,
+                    source_text=fact,
+                    recorded_at=getattr(edge, "reference_time", None)
+                    or getattr(edge, "created_at", None)
+                    or datetime.now(timezone.utc),
+                    confidence=1.0,
+                )
+                for episode_id in getattr(edge, "episodes", [])[:3]
+            ]
+            if not evidence:
+                evidence = [
+                    Evidence(
+                        evidence_id=getattr(edge, "uuid", "graphiti_edge"),
+                        source_text=fact,
+                        recorded_at=getattr(edge, "created_at", None)
+                        or datetime.now(timezone.utc),
+                    )
+                ]
+            items.append(
+                RetrievalItem(
+                    item_id=getattr(edge, "uuid", fact),
+                    kind="fact",
+                    title=getattr(edge, "name", "Graphiti fact"),
+                    matched_text=fact,
+                    score=1.0,
+                    why_matched="Matched by Graphiti hybrid graph search.",
+                    evidence=evidence,
+                    sensitivity=[SensitivityLabel.SECONDHAND]
+                    if _looks_secondhand(fact)
+                    else [],
+                    is_secondhand=_looks_secondhand(fact),
+                )
+            )
+        semantic_items = self._semantic_projection_search(
+            query,
+            limit=max(limit * 4, 20),
+            include_sensitive=include_sensitive,
+        )
+        local_items = super().search(
+            query,
+            limit=max(limit * 4, 20),
+            include_sensitive=include_sensitive,
+            mode=mode,
+        )
+        candidates = _apply_intent_reranking(query, [*semantic_items, *local_items, *items])
+        return _merge_retrieval_items(candidates, limit=limit)
+
+    async def _add_episode(self, interaction: SocialInteraction) -> Any:
+        from graphiti_core.nodes import EpisodeType
+
+        return await self._graphiti.add_episode(
+            name=f"social_interaction:{interaction.occurred_at or datetime.now(timezone.utc)}",
+            episode_body=format_graphiti_episode(interaction),
+            source_description="people-network-memory record_interaction",
+            reference_time=interaction.occurred_at or datetime.now(timezone.utc),
+            source=EpisodeType.text,
+            entity_types=ENTITY_TYPES,
+            edge_types=EDGE_TYPES,
+            custom_extraction_instructions=CUSTOM_EXTRACTION_INSTRUCTIONS,
+        )
+
+    def _semantic_projection_search(
+        self,
+        query: str,
+        *,
+        limit: int,
+        include_sensitive: bool,
+    ) -> list[RetrievalItem]:
+        path = semantic_index_path(self._config)
+        if not path.exists():
+            return []
+        try:
+            settings = EmbeddingSettings.from_config(self._config)
+            client = OpenAICompatibleEmbeddingClient(settings)
+            return SemanticProjectionIndex(path).search(
+                query,
+                embed_texts=client.embed,
+                limit=limit,
+                include_sensitive=include_sensitive,
+            )
+        except Exception:
+            return []
+
+
+def semantic_index_path(config: PeopleMemoryConfig) -> Path:
+    return Path(config.data_path).expanduser() / "people-memory.semantic-cache.json"
+
+
+def _build_graphiti(config: PeopleMemoryConfig) -> Any:
+    if not config.llm_provider or not config.llm_model:
+        raise BackendUnavailableError(
+            "Graphiti requires PEOPLE_MEMORY_LLM_PROVIDER and PEOPLE_MEMORY_LLM_MODEL."
+        )
+    from graphiti_core import Graphiti
+    from graphiti_core.cross_encoder.client import CrossEncoderClient
+    from graphiti_core.driver.kuzu_driver import KuzuDriver
+    from graphiti_core.embedder.openai import OpenAIEmbedder, OpenAIEmbedderConfig
+    from graphiti_core.llm_client.config import LLMConfig
+    from graphiti_core.llm_client.openai_generic_client import OpenAIGenericClient
+
+    embedding = EmbeddingSettings.from_config(config)
+    llm_config = LLMConfig(
+        api_key=config.llm_api_key or "local",
+        model=config.llm_model,
+        small_model=config.llm_model,
+        base_url=config.llm_base_url,
+    )
+    llm_client = _build_llm_client(config, llm_config, OpenAIGenericClient)
+    embedder = OpenAIEmbedder(
+        config=OpenAIEmbedderConfig(
+            api_key=embedding.api_key,
+            embedding_model=embedding.model,
+            embedding_dim=embedding.dimension or 1024,
+            base_url=embedding.base_url,
+        )
+    )
+    cross_encoder = _build_cross_encoder(CrossEncoderClient)
+    if config.graph_backend_kind == "kuzu":
+        kuzu_path = Path(config.graphiti_kuzu_path).expanduser()
+        kuzu_path.parent.mkdir(parents=True, exist_ok=True)
+        driver = KuzuDriver(db=str(kuzu_path))
+        graphiti = Graphiti(
+            graph_driver=driver,
+            llm_client=llm_client,
+            embedder=embedder,
+            cross_encoder=cross_encoder,
+        )
+        _run_async(_ensure_kuzu_fulltext_indices(driver))
+        return graphiti
+    raise BackendUnavailableError(
+        f"Graphiti backend {config.graph_backend_kind} is not wired in this adapter yet."
+    )
+
+
+def _build_llm_client(
+    config: PeopleMemoryConfig, llm_config: Any, openai_generic_client: Any
+) -> Any:
+    if config.llm_response_format == "json_schema":
+        return openai_generic_client(config=llm_config)
+
+    class JsonObjectCompatibleClient(openai_generic_client):  # type: ignore[misc, valid-type]
+        async def _generate_response(
+            self,
+            messages: list[Any],
+            response_model: type[Any] | None = None,
+            max_tokens: int = 16384,
+            model_size: Any = None,
+        ) -> dict[str, Any]:
+            import openai
+            from graphiti_core.llm_client.errors import RateLimitError
+
+            if response_model is not None and messages:
+                serialized_model = json.dumps(response_model.model_json_schema())
+                messages[-1].content += (
+                    "\n\nRespond with exactly one JSON object matching this JSON Schema:\n"
+                    f"{serialized_model}"
+                )
+            openai_messages: list[dict[str, str]] = []
+            for message in messages:
+                message.content = self._clean_input(message.content)
+                if message.role in {"user", "system"}:
+                    openai_messages.append(
+                        {"role": message.role, "content": message.content}
+                    )
+            try:
+                request: dict[str, Any] = {
+                    "model": self.model,
+                    "messages": openai_messages,
+                    "temperature": self.temperature,
+                    "max_tokens": max_tokens or self.max_tokens,
+                }
+                if config.llm_response_format == "json_object":
+                    request["response_format"] = {"type": "json_object"}
+                response = await self.client.chat.completions.create(**request)
+                result = response.choices[0].message.content or "{}"
+                parsed = _parse_json_object(result)
+                if response_model is not None and isinstance(parsed, dict):
+                    valid_fields = set(response_model.model_fields)
+                    parsed = {
+                        key: value
+                        for key, value in parsed.items()
+                        if key in valid_fields
+                    }
+                return parsed
+            except openai.RateLimitError as exc:
+                raise RateLimitError from exc
+
+    return JsonObjectCompatibleClient(config=llm_config)
+
+
+def _build_cross_encoder(cross_encoder_client: Any) -> Any:
+    class PassThroughCrossEncoder(cross_encoder_client):  # type: ignore[misc, valid-type]
+        async def rank(self, query: str, passages: list[str]) -> list[tuple[str, float]]:
+            query_terms = {term.lower() for term in query.split() if term.strip()}
+
+            def score(passage: str) -> float:
+                lowered = passage.lower()
+                return float(sum(1 for term in query_terms if term in lowered))
+
+            return sorted(
+                ((passage, score(passage)) for passage in passages),
+                key=lambda item: item[1],
+                reverse=True,
+            )
+
+    return PassThroughCrossEncoder()
+
+
+async def _ensure_kuzu_fulltext_indices(driver: Any) -> None:
+    from graphiti_core.driver.driver import GraphProvider
+    from graphiti_core.graph_queries import get_fulltext_indices
+
+    kuzu_logger = logging.getLogger("graphiti_core.driver.kuzu_driver")
+    duplicate_index_filter = _KuzuDuplicateIndexLogFilter()
+    kuzu_logger.addFilter(duplicate_index_filter)
+    try:
+        for query in get_fulltext_indices(GraphProvider.KUZU):
+            try:
+                await driver.execute_query(query)
+            except Exception as exc:
+                if not _is_duplicate_kuzu_fulltext_index_exception(exc):
+                    raise
+    finally:
+        kuzu_logger.removeFilter(duplicate_index_filter)
+
+
+class _KuzuDuplicateIndexLogFilter(logging.Filter):
+    def filter(self, record: logging.LogRecord) -> bool:
+        return not _is_duplicate_kuzu_fulltext_index_log(record.getMessage())
+
+
+def _is_duplicate_kuzu_fulltext_index_log(message: str) -> bool:
+    lowered = message.lower()
+    return (
+        "error executing kuzu query" in lowered
+        and "create_fts_index" in lowered
+        and "already exists" in lowered
+        and "index" in lowered
+    )
+
+
+def _is_duplicate_kuzu_fulltext_index_exception(exc: Exception) -> bool:
+    lowered = str(exc).lower()
+    return "already exists" in lowered and "index" in lowered
+
+
+def _run_async(awaitable: Any) -> Any:
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        return asyncio.run(awaitable)
+    raise RuntimeError("GraphitiGraphStore cannot be called from an active event loop yet.")
+
+
+def _with_retries(operation: Any, *, label: str, attempts: int = 3) -> Any:
+    last_exc: Exception | None = None
+    for attempt in range(1, attempts + 1):
+        try:
+            return operation()
+        except Exception as exc:  # pragma: no cover - retry shape is tested with fakes
+            last_exc = exc
+            if attempt >= attempts or not _is_retryable_graphiti_error(exc):
+                raise
+            time.sleep(min(2 ** (attempt - 1), 4))
+    raise RuntimeError(f"{label} failed without an exception") from last_exc
+
+
+def _is_retryable_graphiti_error(exc: Exception) -> bool:
+    lowered = str(exc).lower()
+    retryable_terms = [
+        "500",
+        "502",
+        "503",
+        "504",
+        "internalserviceerror",
+        "internal server error",
+        "timeout",
+        "timed out",
+        "temporarily unavailable",
+        "rate limit",
+        "no valid json object found",
+    ]
+    return any(term in lowered for term in retryable_terms)
+
+
+def _parse_json_object(text: str) -> dict[str, Any]:
+    stripped = _strip_json_fence(text.strip())
+    parsed = _parse_jsonish_object(stripped)
+    if not isinstance(parsed, dict):
+        raise json.JSONDecodeError("Expected a JSON object", stripped, 0)
+    return parsed
+
+
+def _strip_json_fence(text: str) -> str:
+    if not text.startswith("```"):
+        return text
+    lines = text.splitlines()
+    if len(lines) >= 3 and lines[-1].strip() == "```":
+        return "\n".join(lines[1:-1]).strip()
+    return text
+
+
+def _first_balanced_json_object(text: str) -> str:
+    start = text.find("{")
+    if start < 0:
+        raise json.JSONDecodeError("No JSON object found", text, 0)
+    depth = 0
+    in_string = False
+    escaped = False
+    for index in range(start, len(text)):
+        char = text[index]
+        if escaped:
+            escaped = False
+            continue
+        if char == "\\" and in_string:
+            escaped = True
+            continue
+        if char == '"':
+            in_string = not in_string
+            continue
+        if in_string:
+            continue
+        if char == "{":
+            depth += 1
+        elif char == "}":
+            depth -= 1
+            if depth == 0:
+                return text[start : index + 1]
+    raise json.JSONDecodeError("Unterminated JSON object", text, start)
+
+
+def _parse_jsonish_object(text: str) -> Any:
+    candidates = [text]
+    try:
+        balanced = _first_balanced_json_object(text)
+    except json.JSONDecodeError:
+        balanced = None
+    if balanced and balanced != text:
+        candidates.append(balanced)
+    for candidate in candidates:
+        try:
+            return json.loads(candidate)
+        except json.JSONDecodeError:
+            pass
+        try:
+            return ast.literal_eval(candidate)
+        except (ValueError, SyntaxError):
+            pass
+        try:
+            return json.loads(_repair_common_json(candidate))
+        except json.JSONDecodeError:
+            pass
+    raise json.JSONDecodeError("No valid JSON object found", text, 0)
+
+
+def _repair_common_json(text: str) -> str:
+    repaired = re.sub(r",\s*([}\]])", r"\1", text)
+    repaired = re.sub(r"([{,]\s*)([A-Za-z_][A-Za-z0-9_]*)\s*:", r'\1"\2":', repaired)
+    repaired = re.sub(r"\bTrue\b", "true", repaired)
+    repaired = re.sub(r"\bFalse\b", "false", repaired)
+    repaired = re.sub(r"\bNone\b", "null", repaired)
+    return repaired
+
+
+class _AsyncLoopRunner:
+    def __init__(self) -> None:
+        self._loop = asyncio.new_event_loop()
+        self._ready = threading.Event()
+        self._thread = threading.Thread(target=self._run_loop, daemon=True)
+        self._thread.start()
+        self._ready.wait(timeout=5)
+
+    def run(self, awaitable: Any, *, timeout_seconds: float | None = None) -> Any:
+        if self._loop.is_closed():
+            raise RuntimeError("Graphiti async loop is closed")
+        future = asyncio.run_coroutine_threadsafe(awaitable, self._loop)
+        try:
+            return future.result(timeout=timeout_seconds)
+        except FutureTimeoutError as exc:
+            future.cancel()
+            raise TimeoutError("Graphiti async operation timed out") from exc
+
+    def close(self) -> None:
+        if self._loop.is_closed():
+            return
+        self._loop.call_soon_threadsafe(self._loop.stop)
+        self._thread.join(timeout=2)
+        self._loop.close()
+
+    def _run_loop(self) -> None:
+        asyncio.set_event_loop(self._loop)
+        self._ready.set()
+        self._loop.run_forever()
+
+    def __del__(self) -> None:  # pragma: no cover - best-effort cleanup
+        try:
+            self.close()
+        except Exception:
+            pass
+
+
+def _looks_sensitive(text: str) -> bool:
+    lowered = text.lower()
+    return any(term in lowered for term in ["private", "sensitive", "confidential", "compensation"])
+
+
+def _looks_secondhand(text: str) -> bool:
+    lowered = text.lower()
+    return any(term in lowered for term in [" said ", " mentioned ", "told", "提到"])
+
+
+def _identity_map_for_cache(
+    store: GraphitiGraphStore, interaction: SocialInteraction
+) -> dict[str, str | None]:
+    identity_map: dict[str, str | None] = {}
+    for ref in _iter_interaction_refs(interaction):
+        key = _ref_key(ref)
+        if key in identity_map:
+            continue
+        if ref.person_id:
+            identity_map[key] = ref.person_id
+            continue
+        candidates = store.find_identity_candidates(ref)
+        identity_map[key] = candidates[0].person_id if candidates and candidates[0].score >= 0.9 else None
+    return identity_map
+
+
+def _iter_interaction_refs(interaction: SocialInteraction) -> list[PersonRef]:
+    refs: list[PersonRef] = []
+    refs.extend(participant.person for participant in interaction.participants)
+    refs.extend(mentioned.person for mentioned in interaction.mentioned_people)
+    refs.extend(mentioned.mentioned_by for mentioned in interaction.mentioned_people if mentioned.mentioned_by)
+    for claim in interaction.attributed_claims:
+        if claim.speaker:
+            refs.append(claim.speaker)
+        if claim.subject:
+            refs.append(claim.subject)
+    refs.extend(fact.subject for fact in interaction.direct_facts)
+    for follow_up in interaction.follow_ups:
+        refs.extend(follow_up.related_people)
+    for relationship in interaction.relationships:
+        refs.extend([relationship.source, relationship.target])
+    return refs
+
+
+def _ref_key(ref: PersonRef) -> str:
+    return ref.person_id or ref.email or ref.phone or ref.label
+
+
+def _apply_intent_reranking(query: str, items: list[RetrievalItem]) -> list[RetrievalItem]:
+    mention_target = mentioned_query_target(query)
+    if mention_target:
+        mentioned_items: list[RetrievalItem] = []
+        for item in items:
+            text = f"{item.title} {item.matched_text}"
+            if not text_answers_mentioned_query(text, mention_target):
+                continue
+            if _title_represents_person(item.title, mention_target):
+                continue
+            mentioned_items.append(item.model_copy(update={"score": item.score + 1.5}))
+        if mentioned_items:
+            return mentioned_items
+
+    if is_follow_up_query(query):
+        follow_up_items = [
+            item
+            for item in items
+            if item.kind == "follow_up"
+            or "follow up" in f"{item.title} {item.matched_text}".lower()
+            or "follow-up" in f"{item.title} {item.matched_text}".lower()
+        ]
+        if follow_up_items:
+            return [
+                item.model_copy(update={"score": item.score + 2.0})
+                for item in follow_up_items
+            ]
+    return items
+
+
+def _title_represents_person(title: str, person_label: str) -> bool:
+    lowered_title = title.strip().lower()
+    lowered_person = person_label.strip().lower()
+    return lowered_title in {
+        lowered_person,
+        f"interaction with {lowered_person}",
+        f"claim involving {lowered_person}",
+        f"relationship involving {lowered_person}",
+        f"follow-up for {lowered_person}",
+    }
+
+
+def _merge_retrieval_items(items: list[RetrievalItem], *, limit: int) -> list[RetrievalItem]:
+    merged: dict[tuple[str, str, str], RetrievalItem] = {}
+    for item in items:
+        key = (item.kind, item.title, item.matched_text)
+        existing = merged.get(key)
+        if existing is None or item.score > existing.score:
+            merged[key] = item
+    return sorted(merged.values(), key=lambda item: item.score, reverse=True)[:limit]