apsimo 1.3.0__py3-none-any.whl

apsimo/world_model/extraction/formats/csv_fmt.py

@@ -0,0 +1,68 @@
+"""CSV entity extractor."""
+
+from __future__ import annotations
+
+import csv
+import io
+import logging
+from typing import Any, Dict, List
+
+from apsimo.world_model.extraction.base import FormatExtractor, ExtractedEntity
+
+logger = logging.getLogger(__name__)
+
+
+class CSVExtractor(FormatExtractor):
+    """Extract entities from CSV data.
+
+    Each row becomes an entity. Column headers are used as attribute keys.
+    A 'name' or 'title' column (if present) is used as the entity name.
+    """
+
+    def supported_formats(self) -> List[str]:
+        return ["text/csv"]
+
+    async def extract_text(self, content: bytes, metadata: Dict[str, Any]) -> str:
+        try:
+            text = content.decode("utf-8")
+        except UnicodeDecodeError:
+            try:
+                text = content.decode("latin-1")
+            except Exception:
+                return ""
+        return text
+
+    async def extract_entities(self, content: bytes, metadata: Dict[str, Any]) -> List[ExtractedEntity]:
+        try:
+            text = content.decode("utf-8")
+        except UnicodeDecodeError:
+            try:
+                text = content.decode("latin-1")
+            except Exception as e:
+                logger.warning("Failed to decode CSV: %s", e)
+                return []
+
+        reader = csv.DictReader(io.StringIO(text))
+        if not reader.fieldnames:
+            return []
+
+        entities = []
+        for i, row in enumerate(reader):
+            # Find the name column
+            name = row.get("name") or row.get("title") or row.get("Name") or row.get("Title") or f"row-{i}"
+            etype = row.get("type") or row.get("Type") or "entity"
+
+            # All other columns become attributes
+            attrs = {}
+            for k, v in row.items():
+                if k and k.lower() not in ("name", "title", "type") and v:
+                    attrs[k] = v
+
+            entities.append(ExtractedEntity(
+                name=str(name).strip(),
+                entity_type=str(etype).strip(),
+                attributes=attrs,
+                confidence=0.85,
+            ))
+
+        return entities

apsimo/world_model/extraction/formats/html_fmt.py

@@ -0,0 +1,72 @@
+"""HTML entity extractor (requires BeautifulSoup)."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List
+
+from apsimo.world_model.extraction.base import FormatExtractor, ExtractedEntity
+
+logger = logging.getLogger(__name__)
+
+try:
+    from bs4 import BeautifulSoup
+    HAS_BS4 = True
+except ImportError:
+    HAS_BS4 = False
+
+
+class HTMLExtractor(FormatExtractor):
+    """Extract text and entities from HTML via BeautifulSoup."""
+
+    def supported_formats(self) -> List[str]:
+        return ["text/html"]
+
+    async def extract_text(self, content: bytes, metadata: Dict[str, Any]) -> str:
+        if not HAS_BS4:
+            logger.warning("BeautifulSoup not installed — cannot extract HTML text")
+            return ""
+
+        try:
+            soup = BeautifulSoup(content, "html.parser")
+            # Remove script and style elements
+            for tag in soup(["script", "style", "nav", "footer", "header"]):
+                tag.decompose()
+            return soup.get_text(separator="\n", strip=True)
+        except Exception as e:
+            logger.warning("Failed to extract HTML text: %s", e)
+            return ""
+
+    async def extract_entities(self, content: bytes, metadata: Dict[str, Any]) -> List[ExtractedEntity]:
+        if not HAS_BS4:
+            return []
+
+        try:
+            soup = BeautifulSoup(content, "html.parser")
+            entities = []
+
+            # Extract metadata entities
+            title = soup.find("title")
+            if title and title.string:
+                entities.append(ExtractedEntity(
+                    name=title.string.strip(),
+                    entity_type="webpage",
+                    attributes={"type": "title"},
+                    confidence=0.95,
+                ))
+
+            # Extract links as relationship hints
+            for a in soup.find_all("a", href=True)[:20]:
+                link_text = a.get_text(strip=True)
+                if link_text:
+                    entities.append(ExtractedEntity(
+                        name=link_text,
+                        entity_type="link",
+                        attributes={"href": a["href"]},
+                        confidence=0.7,
+                    ))
+
+            return entities
+        except Exception as e:
+            logger.warning("Failed to extract HTML entities: %s", e)
+            return []

apsimo/world_model/extraction/formats/json_fmt.py

@@ -0,0 +1,68 @@
+"""JSON entity extractor."""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any, Dict, List
+
+from apsimo.world_model.extraction.base import FormatExtractor, ExtractedEntity
+
+logger = logging.getLogger(__name__)
+
+
+class JSONExtractor(FormatExtractor):
+    """Extract entities from structured JSON data.
+
+    Handles two patterns:
+    1. Array of objects: each object becomes an entity
+    2. Nested objects: top-level keys become entity types
+    """
+
+    def supported_formats(self) -> List[str]:
+        return ["application/json"]
+
+    async def extract_text(self, content: bytes, metadata: Dict[str, Any]) -> str:
+        try:
+            data = json.loads(content)
+            return json.dumps(data, indent=2, ensure_ascii=False)
+        except (json.JSONDecodeError, UnicodeDecodeError):
+            return ""
+
+    async def extract_entities(self, content: bytes, metadata: Dict[str, Any]) -> List[ExtractedEntity]:
+        try:
+            data = json.loads(content)
+        except (json.JSONDecodeError, UnicodeDecodeError) as e:
+            logger.warning("Failed to parse JSON: %s", e)
+            return []
+
+        entities = []
+
+        # Array of objects — each becomes an entity
+        if isinstance(data, list):
+            for i, item in enumerate(data):
+                if isinstance(item, dict):
+                    name = item.get("name") or item.get("title") or item.get("id") or f"item-{i}"
+                    etype = item.get("type", "entity")
+                    attrs = {k: v for k, v in item.items() if k not in ("name", "title", "id", "type")}
+                    entities.append(ExtractedEntity(
+                        name=str(name),
+                        entity_type=str(etype),
+                        attributes=attrs,
+                        confidence=0.9,
+                    ))
+
+        # Single object — top-level keys become attributes
+        elif isinstance(data, dict):
+            # If it has a 'name' or 'title', treat as one entity
+            name = data.get("name") or data.get("title") or metadata.get("filename", "json-object")
+            etype = data.get("type", "entity")
+            attrs = {k: v for k, v in data.items() if k not in ("name", "title", "type") and isinstance(v, (str, int, float, bool))}
+            entities.append(ExtractedEntity(
+                name=str(name),
+                entity_type=str(etype),
+                attributes=attrs,
+                confidence=0.9,
+            ))
+
+        return entities

apsimo/world_model/extraction/formats/pdf.py

@@ -0,0 +1,43 @@
+"""PDF entity extractor (requires PyMuPDF)."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List
+
+from apsimo.world_model.extraction.base import FormatExtractor, ExtractedEntity
+
+logger = logging.getLogger(__name__)
+
+try:
+    import fitz  # PyMuPDF
+    HAS_PYMUPDF = True
+except ImportError:
+    HAS_PYMUPDF = False
+
+
+class PDFExtractor(FormatExtractor):
+    """Extract text and entities from PDF files via PyMuPDF."""
+
+    def supported_formats(self) -> List[str]:
+        return ["application/pdf"]
+
+    async def extract_text(self, content: bytes, metadata: Dict[str, Any]) -> str:
+        if not HAS_PYMUPDF:
+            logger.warning("PyMuPDF not installed — cannot extract PDF text")
+            return ""
+
+        try:
+            doc = fitz.open(stream=content, filetype="pdf")
+            text_parts = []
+            for page in doc:
+                text_parts.append(page.get_text())
+            doc.close()
+            return "\n\n".join(text_parts)
+        except Exception as e:
+            logger.warning("Failed to extract PDF text: %s", e)
+            return ""
+
+    async def extract_entities(self, content: bytes, metadata: Dict[str, Any]) -> List[ExtractedEntity]:
+        # PDFs don't have structured entities — rely on text extraction + LLM
+        return []

apsimo/world_model/extraction/formats/text.py

@@ -0,0 +1,27 @@
+"""Plain text extractor."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List
+
+from apsimo.world_model.extraction.base import FormatExtractor, ExtractedEntity
+
+logger = logging.getLogger(__name__)
+
+
+class TextExtractor(FormatExtractor):
+    """Extract entities from plain text (via LLM)."""
+
+    def supported_formats(self) -> List[str]:
+        return ["text/plain", "text/markdown"]
+
+    async def extract_text(self, content: bytes, metadata: Dict[str, Any]) -> str:
+        try:
+            return content.decode("utf-8")
+        except UnicodeDecodeError:
+            try:
+                return content.decode("latin-1")
+            except Exception:
+                logger.warning("Failed to decode text content")
+                return ""

apsimo/world_model/extraction/llm_extractor.py

@@ -0,0 +1,164 @@
+"""LLM-backed entity extractor — the fallback path for ExtractionPipeline.
+
+``ExtractionPipeline`` first tries each ``FormatExtractor``'s structured
+``extract_entities`` pass. When that returns an empty list (common for
+unstructured PDFs, plain HTML, free-form text) the pipeline falls back to
+``extract_text`` + this LLM extractor.
+
+The LLM is prompted to return a strict JSON array; anything else (prose,
+partial JSON, network errors) degrades to an empty list so the pipeline
+stays non-fatal.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+from typing import Any, Awaitable, Callable, Dict, List, Optional
+
+from apsimo.world_model.extraction.base import ExtractedEntity
+
+logger = logging.getLogger(__name__)
+
+LlmExtractFn = Callable[[str, Dict[str, Any]], Awaitable[List[ExtractedEntity]]]
+
+# Plenty of room for a useful document but bounded so we don't blow tokens.
+MAX_INPUT_CHARS = 12_000
+MAX_RESPONSE_TOKENS = 1_024
+
+_SYSTEM_PROMPT = (
+    "You are an entity extraction engine. Read the document and return a "
+    "JSON array of entities. Each entity must be an object with keys: "
+    '"name" (string, required), "type" (one of: person, company, project, '
+    'product, location, event, concept), "attributes" (object, optional), '
+    '"confidence" (number 0-1, optional). Return ONLY the JSON array: '
+    "no prose, no code fences, no explanation. Return [] if no entities "
+    "are present. Do not invent entities not grounded in the text."
+)
+
+
+def build_llm_extract_fn(
+    llm_router: Any,
+    *,
+    max_input_chars: int = MAX_INPUT_CHARS,
+    max_response_tokens: int = MAX_RESPONSE_TOKENS,
+) -> LlmExtractFn:
+    """Build an ``llm_extract_fn`` bound to the given ``LLMRouter``.
+
+    Parameters
+    ----------
+    llm_router :
+        An object with an ``async complete(messages, ...)`` method that
+        returns an object with a ``.content`` attribute (the ``LLMRouter``
+        from ``apsimo.router.router``).
+    max_input_chars :
+        Truncate the text to this many characters before sending to the LLM.
+    max_response_tokens :
+        Cap on LLM output tokens.
+    """
+
+    async def extract(text: str, metadata: Dict[str, Any]) -> List[ExtractedEntity]:
+        snippet = (text or "").strip()
+        if not snippet:
+            return []
+        if len(snippet) > max_input_chars:
+            snippet = snippet[:max_input_chars]
+        user_prompt = _build_user_prompt(snippet, metadata)
+        try:
+            resp = await llm_router.complete(
+                messages=[
+                    {"role": "system", "content": _SYSTEM_PROMPT},
+                    {"role": "user", "content": user_prompt},
+                ],
+                context={"task": "entity_extraction", "max_tokens": max_response_tokens},
+            )
+        except Exception as exc:
+            logger.warning("LLM entity extraction call failed: %s", exc)
+            return []
+
+        content = getattr(resp, "content", "") or ""
+        return _parse_entity_array(content)
+
+    return extract
+
+
+def _build_user_prompt(text: str, metadata: Dict[str, Any]) -> str:
+    hints: List[str] = []
+    filename = metadata.get("filename") or metadata.get("source")
+    if filename:
+        hints.append(f"Source: {filename}")
+    if metadata.get("mime_type"):
+        hints.append(f"MIME: {metadata['mime_type']}")
+    prefix = ("\n".join(hints) + "\n\n") if hints else ""
+    return f"{prefix}Document:\n---\n{text}\n---\n\nReturn the JSON array."
+
+
+_TYPE_SYNONYMS = {
+    "organization": "company", "org": "company", "business": "company",
+    "place": "location",
+}
+
+
+def _parse_entity_array(raw: str) -> List[ExtractedEntity]:
+    """Best-effort parse of an LLM response into ``ExtractedEntity`` objects.
+
+    Accepts either a bare JSON array or an object with an ``entities`` key
+    (LLMs occasionally wrap the array). Tolerates code fences.
+    """
+    if not raw:
+        return []
+    candidate = _strip_code_fence(raw).strip()
+    try:
+        parsed = json.loads(candidate)
+    except json.JSONDecodeError:
+        # Fall back to the first [...] block in the response.
+        match = re.search(r"\[.*\]", candidate, re.DOTALL)
+        if not match:
+            logger.debug("LLM extractor: no JSON array in response")
+            return []
+        try:
+            parsed = json.loads(match.group(0))
+        except json.JSONDecodeError as exc:
+            logger.debug("LLM extractor: JSON parse failed (%s)", exc)
+            return []
+
+    if isinstance(parsed, dict):
+        parsed = parsed.get("entities") or parsed.get("items") or []
+    if not isinstance(parsed, list):
+        return []
+
+    entities: List[ExtractedEntity] = []
+    for item in parsed:
+        if not isinstance(item, dict):
+            continue
+        name = item.get("name")
+        entity_type = item.get("type") or item.get("entity_type")
+        if not isinstance(name, str) or not isinstance(entity_type, str):
+            continue
+        entity_type = _TYPE_SYNONYMS.get(
+            entity_type.strip().lower(), entity_type.strip().lower())
+        attributes = item.get("attributes")
+        if not isinstance(attributes, dict):
+            attributes = None
+        confidence = item.get("confidence", 0.7)
+        try:
+            confidence = float(confidence)
+        except (TypeError, ValueError):
+            confidence = 0.7
+        entities.append(
+            ExtractedEntity(
+                name=name.strip(),
+                entity_type=entity_type.strip(),
+                attributes=attributes,
+                confidence=max(0.0, min(1.0, confidence)),
+            )
+        )
+    return entities
+
+
+_CODE_FENCE = re.compile(r"^```(?:json)?\s*|\s*```$", re.MULTILINE)
+
+
+def _strip_code_fence(text: str) -> str:
+    return _CODE_FENCE.sub("", text)

apsimo/world_model/extraction/pipeline.py

@@ -0,0 +1,73 @@
+"""Extraction pipeline — orchestrates format detection and entity extraction."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable, Dict, List, Optional
+
+from apsimo.world_model.extraction.base import ExtractedEntity, FormatExtractor
+from apsimo.world_model.extraction.detector import FormatDetector
+
+logger = logging.getLogger(__name__)
+
+
+class ExtractionPipeline:
+    """Orchestrate format detection and entity extraction.
+
+    Flow:
+    1. Detect format
+    2. Try structured extraction (fast, no LLM)
+    3. Fall back to text extraction + LLM
+    """
+
+    def __init__(
+        self,
+        extractors: List[FormatExtractor],
+        llm_extract_fn: Optional[Callable] = None,
+    ):
+        self._extractors: Dict[str, FormatExtractor] = {}
+        for ext in extractors:
+            for fmt in ext.supported_formats():
+                self._extractors[fmt] = ext
+        self._detector = FormatDetector()
+        self._llm_extract_fn = llm_extract_fn
+
+    async def extract(
+        self,
+        content: bytes,
+        filename: str = "",
+        mime_type: str = "",
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> List[ExtractedEntity]:
+        """Extract entities from content.
+
+        1. Detect format
+        2. Try structured extraction (fast, no LLM)
+        3. Fall back to text extraction + LLM
+        """
+        metadata = metadata or {}
+
+        # Detect format
+        fmt = self._detector.detect(content, filename=filename, mime_type=mime_type)
+        logger.debug("Detected format: %s", fmt)
+
+        # Find extractor
+        extractor = self._extractors.get(fmt)
+        if extractor is None:
+            logger.warning("No extractor for format '%s' — skipping", fmt)
+            return []
+
+        # 1. Try structured extraction first
+        entities = await extractor.extract_entities(content, metadata)
+
+        # 2. Fall back to text + LLM
+        if not entities:
+            text = await extractor.extract_text(content, metadata)
+            if text and self._llm_extract_fn:
+                try:
+                    entities = await self._llm_extract_fn(text, metadata)
+                except Exception as e:
+                    logger.warning("LLM extraction failed: %s", e)
+
+        logger.info("Extracted %d entities from %s format", len(entities), fmt)
+        return entities

apsimo/world_model/integrations/__init__.py

@@ -0,0 +1,5 @@
+"""World Model integration bridges."""
+from .mind_model_bridge import MindModelBridge
+from .social_intel_bridge import SocialIntelBridge
+
+__all__ = ["MindModelBridge", "SocialIntelBridge"]

apsimo/world_model/integrations/mind_model_bridge.py

@@ -0,0 +1,115 @@
+"""Mind Model bridge: provides world model context to the mind model pipeline."""
+
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..store import WorldModelStore
+
+
+class MindModelBridge:
+    """Provides world model context to the mind model pipeline."""
+
+    def __init__(self, store: "WorldModelStore") -> None:
+        self._store = store
+
+    async def get_context_for_person(
+        self,
+        person_id: str,
+        as_of: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Fetch world model context relevant to a person's mental state.
+
+        Returns employer, active projects, upcoming events, and shared
+        concepts that may inform state estimation.
+
+        Args:
+            person_id: The world model entity ID of the person.
+            as_of: Optional ISO8601 datetime for temporal queries.
+
+        Returns:
+            Dict with 'employer', 'projects', 'events', 'concepts' keys.
+        """
+        context: Dict[str, Any] = {
+            "employer": None,
+            "projects": [],
+            "events": [],
+            "concepts": [],
+        }
+
+        # Current employer
+        employer_rels = await self._store.query_relationships(
+            source_id=person_id,
+            relationship_type="WM_WORKS_AT",
+            active_only=(as_of is None),
+            min_confidence=0.30,
+            limit=1,
+        )
+        if employer_rels:
+            employer = await self._store.get_entity(employer_rels[0].target_id)
+            context["employer"] = employer
+
+        # Active project memberships
+        project_rels = await self._store.query_relationships(
+            source_id=person_id,
+            relationship_type="WM_MEMBER_OF",
+            target_types=["project"],
+            active_only=(as_of is None),
+            min_confidence=0.30,
+            limit=10,
+        )
+        for rel in project_rels:
+            project = await self._store.get_entity(rel.target_id)
+            if project:
+                context["projects"].append(project)
+
+        # Upcoming events (attended)
+        event_rels = await self._store.query_relationships(
+            source_id=person_id,
+            relationship_type="WM_ATTENDED",
+            target_types=["event"],
+            min_confidence=0.30,
+            limit=10,
+        )
+        for rel in event_rels:
+            event = await self._store.get_entity(rel.target_id)
+            if event:
+                context["events"].append(event)
+
+        # Tagged concepts
+        concept_rels = await self._store.query_relationships(
+            source_id=person_id,
+            relationship_type="WM_TAGGED_WITH",
+            target_types=["concept"],
+            min_confidence=0.30,
+            limit=10,
+        )
+        for rel in concept_rels:
+            concept = await self._store.get_entity(rel.target_id)
+            if concept:
+                context["concepts"].append(concept)
+
+        return context
+
+    async def enrich_mind_signal(
+        self,
+        person_id: str,
+        signal_text: str,
+    ) -> Optional[str]:
+        """Cross-reference a text signal with world model entities.
+
+        Returns the entity ID if a matching EventEntity or ConceptEntity
+        is found, otherwise None.
+        """
+        if not signal_text:
+            return None
+
+        # Search events and concepts
+        candidates = await self._store.find_entities(
+            query=signal_text,
+            min_confidence=0.30,
+            limit=5,
+        )
+        for candidate in candidates:
+            if candidate.entity_type in ("event", "concept"):
+                return candidate.id
+        return None