alma-memory 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

alma/graph/extraction.py

@@ -1,198 +1,198 @@
-"""
-ALMA Graph Memory - Entity Extraction.
-
-LLM-powered extraction of entities and relationships from text.
-"""
-
-import logging
-from dataclasses import dataclass
-from typing import Dict, List, Optional, Tuple
-
-from alma.graph.store import Entity, Relationship
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class ExtractionConfig:
-    """Configuration for entity extraction."""
-
-    provider: str = "openai"
-    model: str = "gpt-4o-mini"
-    temperature: float = 0.1
-
-
-class EntityExtractor:
-    """
-    LLM-powered entity and relationship extraction.
-
-    Extracts entities (people, organizations, tools, concepts) and
-    relationships between them from text.
-    """
-
-    EXTRACTION_PROMPT = """Extract entities and relationships from the following text.
-
-Entities are things like:
-- People (names, roles)
-- Organizations (companies, teams)
-- Tools/Technologies (software, frameworks)
-- Concepts (methodologies, patterns)
-- Locations (places, regions)
-
-Relationships describe how entities are connected:
-- WORKS_AT (person -> organization)
-- USES (entity -> tool)
-- KNOWS (person -> person)
-- CREATED_BY (thing -> person/org)
-- PART_OF (entity -> larger entity)
-- RELATED_TO (general relationship)
-
-TEXT:
-{text}
-
-Respond in JSON format:
-```json
-{{
-  "entities": [
-    {{"id": "unique-id", "name": "Entity Name", "type": "person|organization|tool|concept|location"}}
-  ],
-  "relationships": [
-    {{"source": "entity-id", "target": "entity-id", "type": "RELATIONSHIP_TYPE", "properties": {{}}}}
-  ]
-}}
-```
-
-Only extract entities and relationships that are explicitly mentioned or strongly implied.
-"""
-
-    def __init__(self, config: Optional[ExtractionConfig] = None):
-        self.config = config or ExtractionConfig()
-        self._client = None
-
-    def _get_client(self):
-        """Lazy initialization of LLM client."""
-        if self._client is None:
-            if self.config.provider == "openai":
-                from openai import OpenAI
-
-                self._client = OpenAI()
-            elif self.config.provider == "anthropic":
-                from anthropic import Anthropic
-
-                self._client = Anthropic()
-        return self._client
-
-    def extract(
-        self,
-        text: str,
-        existing_entities: Optional[List[Entity]] = None,
-    ) -> Tuple[List[Entity], List[Relationship]]:
-        """
-        Extract entities and relationships from text.
-
-        Args:
-            text: Text to extract from
-            existing_entities: Optional list of known entities for linking
-
-        Returns:
-            Tuple of (entities, relationships)
-        """
-        import json
-        import re
-        import uuid
-
-        prompt = self.EXTRACTION_PROMPT.format(text=text)
-
-        client = self._get_client()
-
-        if self.config.provider == "openai":
-            response = client.chat.completions.create(
-                model=self.config.model,
-                messages=[{"role": "user", "content": prompt}],
-                temperature=self.config.temperature,
-            )
-            raw_response = response.choices[0].message.content
-        elif self.config.provider == "anthropic":
-            response = client.messages.create(
-                model=self.config.model,
-                max_tokens=2000,
-                messages=[{"role": "user", "content": prompt}],
-            )
-            raw_response = response.content[0].text
-        else:
-            raise ValueError(f"Unsupported provider: {self.config.provider}")
-
-        # Parse JSON from response
-        json_match = re.search(r"```json\s*(.*?)\s*```", raw_response, re.DOTALL)
-        if json_match:
-            json_str = json_match.group(1)
-        else:
-            json_match = re.search(r"\{.*\}", raw_response, re.DOTALL)
-            json_str = json_match.group(0) if json_match else "{}"
-
-        try:
-            data = json.loads(json_str)
-        except json.JSONDecodeError:
-            logger.warning(f"Failed to parse extraction response: {raw_response[:200]}")
-            return [], []
-
-        # Build entities
-        entities = []
-        entity_map = {}  # id -> Entity
-
-        for e in data.get("entities", []):
-            entity_id = e.get("id") or str(uuid.uuid4())[:8]
-            entity = Entity(
-                id=entity_id,
-                name=e.get("name", "Unknown"),
-                entity_type=e.get("type", "concept"),
-            )
-            entities.append(entity)
-            entity_map[entity_id] = entity
-
-            # Also map by name for relationship linking
-            entity_map[e.get("name", "").lower()] = entity
-
-        # Build relationships
-        relationships = []
-        for r in data.get("relationships", []):
-            source_id = r.get("source")
-            target_id = r.get("target")
-
-            # Try to resolve IDs
-            source = entity_map.get(source_id) or entity_map.get(
-                source_id.lower() if source_id else ""
-            )
-            target = entity_map.get(target_id) or entity_map.get(
-                target_id.lower() if target_id else ""
-            )
-
-            if source and target:
-                rel = Relationship(
-                    id=f"{source.id}-{r.get('type', 'RELATED')}-{target.id}",
-                    source_id=source.id,
-                    target_id=target.id,
-                    relation_type=r.get("type", "RELATED_TO"),
-                    properties=r.get("properties", {}),
-                )
-                relationships.append(rel)
-
-        return entities, relationships
-
-    def extract_from_conversation(
-        self,
-        messages: List[Dict[str, str]],
-    ) -> Tuple[List[Entity], List[Relationship]]:
-        """
-        Extract entities and relationships from a conversation.
-
-        Args:
-            messages: List of {"role": "...", "content": "..."}
-
-        Returns:
-            Tuple of (entities, relationships)
-        """
-        # Combine messages into text
-        text = "\n".join(f"{msg['role'].upper()}: {msg['content']}" for msg in messages)
-
-        return self.extract(text)
+"""
+ALMA Graph Memory - Entity Extraction.
+
+LLM-powered extraction of entities and relationships from text.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+from alma.graph.store import Entity, Relationship
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ExtractionConfig:
+    """Configuration for entity extraction."""
+
+    provider: str = "openai"
+    model: str = "gpt-4o-mini"
+    temperature: float = 0.1
+
+
+class EntityExtractor:
+    """
+    LLM-powered entity and relationship extraction.
+
+    Extracts entities (people, organizations, tools, concepts) and
+    relationships between them from text.
+    """
+
+    EXTRACTION_PROMPT = """Extract entities and relationships from the following text.
+
+Entities are things like:
+- People (names, roles)
+- Organizations (companies, teams)
+- Tools/Technologies (software, frameworks)
+- Concepts (methodologies, patterns)
+- Locations (places, regions)
+
+Relationships describe how entities are connected:
+- WORKS_AT (person -> organization)
+- USES (entity -> tool)
+- KNOWS (person -> person)
+- CREATED_BY (thing -> person/org)
+- PART_OF (entity -> larger entity)
+- RELATED_TO (general relationship)
+
+TEXT:
+{text}
+
+Respond in JSON format:
+```json
+{{
+  "entities": [
+    {{"id": "unique-id", "name": "Entity Name", "type": "person|organization|tool|concept|location"}}
+  ],
+  "relationships": [
+    {{"source": "entity-id", "target": "entity-id", "type": "RELATIONSHIP_TYPE", "properties": {{}}}}
+  ]
+}}
+```
+
+Only extract entities and relationships that are explicitly mentioned or strongly implied.
+"""
+
+    def __init__(self, config: Optional[ExtractionConfig] = None):
+        self.config = config or ExtractionConfig()
+        self._client = None
+
+    def _get_client(self):
+        """Lazy initialization of LLM client."""
+        if self._client is None:
+            if self.config.provider == "openai":
+                from openai import OpenAI
+
+                self._client = OpenAI()
+            elif self.config.provider == "anthropic":
+                from anthropic import Anthropic
+
+                self._client = Anthropic()
+        return self._client
+
+    def extract(
+        self,
+        text: str,
+        existing_entities: Optional[List[Entity]] = None,
+    ) -> Tuple[List[Entity], List[Relationship]]:
+        """
+        Extract entities and relationships from text.
+
+        Args:
+            text: Text to extract from
+            existing_entities: Optional list of known entities for linking
+
+        Returns:
+            Tuple of (entities, relationships)
+        """
+        import json
+        import re
+        import uuid
+
+        prompt = self.EXTRACTION_PROMPT.format(text=text)
+
+        client = self._get_client()
+
+        if self.config.provider == "openai":
+            response = client.chat.completions.create(
+                model=self.config.model,
+                messages=[{"role": "user", "content": prompt}],
+                temperature=self.config.temperature,
+            )
+            raw_response = response.choices[0].message.content
+        elif self.config.provider == "anthropic":
+            response = client.messages.create(
+                model=self.config.model,
+                max_tokens=2000,
+                messages=[{"role": "user", "content": prompt}],
+            )
+            raw_response = response.content[0].text
+        else:
+            raise ValueError(f"Unsupported provider: {self.config.provider}")
+
+        # Parse JSON from response
+        json_match = re.search(r"```json\s*(.*?)\s*```", raw_response, re.DOTALL)
+        if json_match:
+            json_str = json_match.group(1)
+        else:
+            json_match = re.search(r"\{.*\}", raw_response, re.DOTALL)
+            json_str = json_match.group(0) if json_match else "{}"
+
+        try:
+            data = json.loads(json_str)
+        except json.JSONDecodeError:
+            logger.warning(f"Failed to parse extraction response: {raw_response[:200]}")
+            return [], []
+
+        # Build entities
+        entities = []
+        entity_map = {}  # id -> Entity
+
+        for e in data.get("entities", []):
+            entity_id = e.get("id") or str(uuid.uuid4())[:8]
+            entity = Entity(
+                id=entity_id,
+                name=e.get("name", "Unknown"),
+                entity_type=e.get("type", "concept"),
+            )
+            entities.append(entity)
+            entity_map[entity_id] = entity
+
+            # Also map by name for relationship linking
+            entity_map[e.get("name", "").lower()] = entity
+
+        # Build relationships
+        relationships = []
+        for r in data.get("relationships", []):
+            source_id = r.get("source")
+            target_id = r.get("target")
+
+            # Try to resolve IDs
+            source = entity_map.get(source_id) or entity_map.get(
+                source_id.lower() if source_id else ""
+            )
+            target = entity_map.get(target_id) or entity_map.get(
+                target_id.lower() if target_id else ""
+            )
+
+            if source and target:
+                rel = Relationship(
+                    id=f"{source.id}-{r.get('type', 'RELATED')}-{target.id}",
+                    source_id=source.id,
+                    target_id=target.id,
+                    relation_type=r.get("type", "RELATED_TO"),
+                    properties=r.get("properties", {}),
+                )
+                relationships.append(rel)
+
+        return entities, relationships
+
+    def extract_from_conversation(
+        self,
+        messages: List[Dict[str, str]],
+    ) -> Tuple[List[Entity], List[Relationship]]:
+        """
+        Extract entities and relationships from a conversation.
+
+        Args:
+            messages: List of {"role": "...", "content": "..."}
+
+        Returns:
+            Tuple of (entities, relationships)
+        """
+        # Combine messages into text
+        text = "\n".join(f"{msg['role'].upper()}: {msg['content']}" for msg in messages)
+
+        return self.extract(text)