Object extractor

trustgraph-flow/pyproject.toml

@@ -78,6 +78,7 @@ graph-embeddings = "trustgraph.embeddings.graph_embeddings:run"
 graph-rag = "trustgraph.retrieval.graph_rag:run"
 kg-extract-agent = "trustgraph.extract.kg.agent:run"
 kg-extract-definitions = "trustgraph.extract.kg.definitions:run"
+kg-extract-objects = "trustgraph.extract.kg.object:run"
 kg-extract-relationships = "trustgraph.extract.kg.relationships:run"
 kg-extract-topics = "trustgraph.extract.kg.topics:run"
 kg-manager = "trustgraph.cores:run"

trustgraph-flow/trustgraph/extract/kg/object/__init__.py

@@ -0,0 +1,3 @@
+
+from . processor import *
+

trustgraph-flow/trustgraph/extract/kg/object/__main__.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-from . extract import run
+from . processor import run
 
 if __name__ == '__main__':
     run()

trustgraph-flow/trustgraph/extract/kg/object/processor.py

@@ -0,0 +1,213 @@
+"""
+Object extraction service - extracts structured objects from text chunks
+based on configured schemas.
+"""
+
+import json
+import logging
+from typing import Dict, List, Any
+
+# Module logger
+logger = logging.getLogger(__name__)
+
+from .... schema import Chunk, ExtractedObject, Metadata
+from .... schema import PromptRequest, PromptResponse
+from .... schema import RowSchema, Field
+
+from .... base import FlowProcessor, ConsumerSpec, ProducerSpec
+from .... base import PromptClientSpec
+
+default_ident = "kg-extract-objects"
+default_concurrency = 1
+
+class Processor(FlowProcessor):
+
+    def __init__(self, **params):
+
+        id = params.get("id")
+        concurrency = params.get("concurrency", 1)
+
+        # Config key for schemas
+        self.config_key = params.get("config_type", "schema")
+
+        super(Processor, self).__init__(
+            **params | {
+                "id": id,
+                "config-type": self.config_key,
+                "concurrency": concurrency,
+            }
+        )
+
+        self.register_specification(
+            ConsumerSpec(
+                name = "input",
+                schema = Chunk,
+                handler = self.on_chunk,
+                concurrency = concurrency,
+            )
+        )
+
+        self.register_specification(
+            PromptClientSpec(
+                request_name = "prompt-request",
+                response_name = "prompt-response",
+            )
+        )
+
+        self.register_specification(
+            ProducerSpec(
+                name = "output", 
+                schema = ExtractedObject
+            )
+        )
+
+        # Register config handler for schema updates
+        self.register_config_handler(self.on_schema_config)
+
+        # Schema storage: name -> RowSchema
+        self.schemas: Dict[str, RowSchema] = {}
+
+    async def on_schema_config(self, config, version):
+        """Handle schema configuration updates"""
+
+        logger.info(f"Loading schema configuration version {version}")
+
+        # Clear existing schemas
+        self.schemas = {}
+
+        # Process each schema in the config
+        for type_key, schema_value in config.items():
+            
+            # Only process schema types
+            if not type_key.startswith(self.config_key + "."):
+                continue
+                
+            # Extract schema name from key (e.g., "schema.customer_records" -> "customer_records")
+            schema_name = type_key[len(self.config_key) + 1:]
+            
+            try:
+                # Parse the JSON schema definition
+                schema_def = json.loads(schema_value)
+                
+                # Create Field objects
+                fields = []
+                for field_def in schema_def.get("fields", []):
+                    field = Field(
+                        name=field_def["name"],
+                        type=field_def["type"],
+                        size=field_def.get("size", 0),
+                        primary=field_def.get("primary_key", False),
+                        description=field_def.get("description", ""),
+                        required=field_def.get("required", False),
+                        enum_values=field_def.get("enum", []),
+                        indexed=field_def.get("indexed", False)
+                    )
+                    fields.append(field)
+                
+                # Create RowSchema
+                row_schema = RowSchema(
+                    name=schema_def.get("name", schema_name),
+                    description=schema_def.get("description", ""),
+                    fields=fields
+                )
+                
+                self.schemas[schema_name] = row_schema
+                logger.info(f"Loaded schema: {schema_name} with {len(fields)} fields")
+                
+            except Exception as e:
+                logger.error(f"Failed to parse schema {schema_name}: {e}", exc_info=True)
+
+        logger.info(f"Schema configuration loaded: {len(self.schemas)} schemas")
+
+    async def extract_objects_for_schema(self, text: str, schema_name: str, schema: RowSchema) -> List[Dict[str, Any]]:
+        """Extract objects from text for a specific schema"""
+        
+        try:
+            # Use prompt client to extract rows based on schema
+            objects = await self.flow("prompt-request").extract_rows(
+                schema=schema,
+                text=text
+            )
+            
+            return objects if isinstance(objects, list) else []
+            
+        except Exception as e:
+            logger.error(f"Failed to extract objects for schema {schema_name}: {e}", exc_info=True)
+            return []
+
+    async def on_chunk(self, msg, consumer, flow):
+        """Process incoming chunk and extract objects"""
+
+        v = msg.value()
+        logger.info(f"Extracting objects from chunk {v.metadata.id}...")
+
+        chunk_text = v.chunk.decode("utf-8")
+
+        # If no schemas configured, log warning and return
+        if not self.schemas:
+            logger.warning("No schemas configured - skipping extraction")
+            return
+
+        try:
+            # Extract objects for each configured schema
+            for schema_name, schema in self.schemas.items():
+                
+                logger.debug(f"Extracting {schema_name} objects from chunk")
+                
+                # Extract objects using prompt
+                objects = await self.extract_objects_for_schema(
+                    chunk_text, 
+                    schema_name, 
+                    schema
+                )
+                
+                # Emit each extracted object
+                for obj in objects:
+                    
+                    # Calculate confidence (could be enhanced with actual confidence from prompt)
+                    confidence = 0.8  # Default confidence
+                    
+                    # Create ExtractedObject
+                    extracted = ExtractedObject(
+                        metadata=Metadata(
+                            id=f"{v.metadata.id}:{schema_name}:{hash(str(obj))}",
+                            metadata=[],
+                            user=v.metadata.user,
+                            collection=v.metadata.collection,
+                        ),
+                        schema_name=schema_name,
+                        values=obj,
+                        confidence=confidence,
+                        source_span=chunk_text[:100]  # First 100 chars as source reference
+                    )
+                    
+                    await flow("output").send(extracted)
+                    logger.debug(f"Emitted extracted object for schema {schema_name}")
+
+        except Exception as e:
+            logger.error(f"Object extraction exception: {e}", exc_info=True)
+
+        logger.debug("Object extraction complete")
+
+    @staticmethod
+    def add_args(parser):
+        """Add command-line arguments"""
+
+        parser.add_argument(
+            '-c', '--concurrency',
+            type=int,
+            default=default_concurrency,
+            help=f'Concurrent processing threads (default: {default_concurrency})'
+        )
+
+        parser.add_argument(
+            '--config-type',
+            default='schema',
+            help='Configuration type prefix for schemas (default: schema)'
+        )
+
+        FlowProcessor.add_args(parser)
+
+def run():
+    """Entry point for kg-extract-objects command"""
+    Processor.launch(default_ident, __doc__)

trustgraph-flow/trustgraph/extract/object/row/__init__.py

@@ -1,3 +0,0 @@
-
-from . extract import *
-

trustgraph-flow/trustgraph/extract/object/row/extract.py

@@ -1,225 +0,0 @@
-
-"""
-Simple decoder, accepts vector+text chunks input, applies analysis to pull
-out a row of fields.  Output as a vector plus object.
-"""
-
-import urllib.parse
-import os
-import logging
-from pulsar.schema import JsonSchema
-
-# Module logger
-logger = logging.getLogger(__name__)
-
-from .... schema import ChunkEmbeddings, Rows, ObjectEmbeddings, Metadata
-from .... schema import RowSchema, Field
-from .... schema import chunk_embeddings_ingest_queue, rows_store_queue
-from .... schema import object_embeddings_store_queue
-from .... schema import prompt_request_queue
-from .... schema import prompt_response_queue
-from .... log_level import LogLevel
-from .... clients.prompt_client import PromptClient
-from .... base import ConsumerProducer
-
-from .... objects.field import Field as FieldParser
-from .... objects.object import Schema
-
-module = ".".join(__name__.split(".")[1:-1])
-
-default_input_queue = chunk_embeddings_ingest_queue
-default_output_queue = rows_store_queue
-default_vector_queue = object_embeddings_store_queue
-default_subscriber = module
-
-class Processor(ConsumerProducer):
-
-    def __init__(self, **params):
-
-        input_queue = params.get("input_queue", default_input_queue)
-        output_queue = params.get("output_queue", default_output_queue)
-        vector_queue = params.get("vector_queue", default_vector_queue)
-        subscriber = params.get("subscriber", default_subscriber)
-        pr_request_queue = params.get(
-            "prompt_request_queue", prompt_request_queue
-        )
-        pr_response_queue = params.get(
-            "prompt_response_queue", prompt_response_queue
-        )
-
-        super(Processor, self).__init__(
-            **params | {
-                "input_queue": input_queue,
-                "output_queue": output_queue,
-                "subscriber": subscriber,
-                "input_schema": ChunkEmbeddings,
-                "output_schema": Rows,
-                "prompt_request_queue": pr_request_queue,
-                "prompt_response_queue": pr_response_queue,
-            }
-        )
-
-        self.vec_prod = self.client.create_producer(
-            topic=vector_queue,
-            schema=JsonSchema(ObjectEmbeddings),
-        )
-
-        __class__.pubsub_metric.info({
-            "input_queue": input_queue,
-            "output_queue": output_queue,
-            "vector_queue": vector_queue,
-            "prompt_request_queue": pr_request_queue,
-            "prompt_response_queue": pr_response_queue,
-            "subscriber": subscriber,
-            "input_schema": ChunkEmbeddings.__name__,
-            "output_schema": Rows.__name__,
-            "vector_schema": ObjectEmbeddings.__name__,
-        })
-
-        flds = __class__.parse_fields(params["field"])
-
-        for fld in flds:
-            logger.debug(f"Field configuration: {fld}")
-
-        self.primary = None
-
-        for f in flds:
-            if f.primary:
-                if self.primary:
-                    raise RuntimeError(
-                        "Only one primary key field is supported"
-                    )
-                self.primary = f
-
-        if self.primary == None:
-            raise RuntimeError(
-                "Must have exactly one primary key field"
-            )
-
-        self.schema = Schema(
-            name = params["name"],
-            description = params["description"],
-            fields = flds
-        )
-
-        self.row_schema=RowSchema(
-            name=self.schema.name,
-            description=self.schema.description,
-            fields=[
-                Field(
-                     name=f.name, type=str(f.type), size=f.size,
-                     primary=f.primary, description=f.description,
-                )
-                for f in self.schema.fields
-            ]
-        )
-
-        self.prompt = PromptClient(
-            pulsar_host=self.pulsar_host,
-            pulsar_api_key=self.pulsar_api_key,
-            input_queue=pr_request_queue,
-            output_queue=pr_response_queue,
-            subscriber = module + "-prompt",
-        )
-
-    @staticmethod
-    def parse_fields(fields):
-        return [ FieldParser.parse(f) for f in fields ]
-
-    def get_rows(self, chunk):
-        return self.prompt.request_rows(self.schema, chunk)
-
-    def emit_rows(self, metadata, rows):
-
-        t = Rows(
-            metadata=metadata, row_schema=self.row_schema, rows=rows
-        )
-        await self.send(t)
-
-    def emit_vec(self, metadata, name, vec, key_name, key):
-
-        r = ObjectEmbeddings(
-            metadata=metadata, vectors=vec, name=name, key_name=key_name, id=key
-        )
-        self.vec_prod.send(r)
-
-    async def handle(self, msg):
-
-        v = msg.value()
-        logger.info(f"Extracting rows from {v.metadata.id}...")
-
-        chunk = v.chunk.decode("utf-8")
-
-        try:
-
-            rows = self.get_rows(chunk)
-
-            self.emit_rows(
-                metadata=v.metadata,
-                rows=rows
-            )
-
-            for row in rows:
-                self.emit_vec(
-                    metadata=v.metadata, vec=v.vectors,
-                    name=self.schema.name, key_name=self.primary.name,
-                    key=row[self.primary.name]
-                )
-
-            for row in rows:
-                logger.debug(f"Extracted row: {row}")
-
-        except Exception as e:
-            logger.error(f"Row extraction exception: {e}", exc_info=True)
-
-        logger.debug("Row extraction complete")
-
-    @staticmethod
-    def add_args(parser):
-
-        ConsumerProducer.add_args(
-            parser, default_input_queue, default_subscriber,
-            default_output_queue,
-        )
-
-        parser.add_argument(
-            '-c', '--vector-queue',
-            default=default_vector_queue,
-            help=f'Vector output queue (default: {default_vector_queue})'
-        )
-
-        parser.add_argument(
-            '--prompt-request-queue',
-            default=prompt_request_queue,
-            help=f'Prompt request queue (default: {prompt_request_queue})',
-        )
-
-        parser.add_argument(
-            '--prompt-response-queue',
-            default=prompt_response_queue,
-            help=f'Prompt response queue (default: {prompt_response_queue})',
-        )
-
-        parser.add_argument(
-            '-f', '--field',
-            required=True,
-            action='append',
-            help=f'Field definition, format name:type:size:pri:descriptionn',
-        )
-
-        parser.add_argument(
-            '-n', '--name',
-            required=True,
-            help=f'Name of row object',
-        )
-
-        parser.add_argument(
-            '-d', '--description',
-            required=True,
-            help=f'Description of object',
-        )
-
-def run():
-
-    Processor.launch(module, __doc__)
-