From 98cefe93dddcfa864be5c72634196237773ed788 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 5 Aug 2025 10:25:51 +0100 Subject: [PATCH] Object extractor --- trustgraph-flow/pyproject.toml | 1 + .../trustgraph/extract/kg/object/__init__.py | 3 + .../{object/row => kg/object}/__main__.py | 2 +- .../trustgraph/extract/kg/object/processor.py | 213 +++++++++++++++++ .../trustgraph/extract/object/__init__.py | 0 .../trustgraph/extract/object/row/__init__.py | 3 - .../trustgraph/extract/object/row/extract.py | 225 ------------------ 7 files changed, 218 insertions(+), 229 deletions(-) create mode 100644 trustgraph-flow/trustgraph/extract/kg/object/__init__.py rename trustgraph-flow/trustgraph/extract/{object/row => kg/object}/__main__.py (69%) create mode 100644 trustgraph-flow/trustgraph/extract/kg/object/processor.py delete mode 100644 trustgraph-flow/trustgraph/extract/object/__init__.py delete mode 100644 trustgraph-flow/trustgraph/extract/object/row/__init__.py delete mode 100755 trustgraph-flow/trustgraph/extract/object/row/extract.py diff --git a/trustgraph-flow/pyproject.toml b/trustgraph-flow/pyproject.toml index 911c91a0..e4abf2f7 100644 --- a/trustgraph-flow/pyproject.toml +++ b/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" diff --git a/trustgraph-flow/trustgraph/extract/kg/object/__init__.py b/trustgraph-flow/trustgraph/extract/kg/object/__init__.py new file mode 100644 index 00000000..9d16af90 --- /dev/null +++ b/trustgraph-flow/trustgraph/extract/kg/object/__init__.py @@ -0,0 +1,3 @@ + +from . processor import * + diff --git a/trustgraph-flow/trustgraph/extract/object/row/__main__.py b/trustgraph-flow/trustgraph/extract/kg/object/__main__.py similarity index 69% rename from trustgraph-flow/trustgraph/extract/object/row/__main__.py rename to trustgraph-flow/trustgraph/extract/kg/object/__main__.py index 403fe672..986c0257 100755 --- a/trustgraph-flow/trustgraph/extract/object/row/__main__.py +++ b/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() diff --git a/trustgraph-flow/trustgraph/extract/kg/object/processor.py b/trustgraph-flow/trustgraph/extract/kg/object/processor.py new file mode 100644 index 00000000..b87cf3ff --- /dev/null +++ b/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__) \ No newline at end of file diff --git a/trustgraph-flow/trustgraph/extract/object/__init__.py b/trustgraph-flow/trustgraph/extract/object/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/trustgraph-flow/trustgraph/extract/object/row/__init__.py b/trustgraph-flow/trustgraph/extract/object/row/__init__.py deleted file mode 100644 index 81287a3c..00000000 --- a/trustgraph-flow/trustgraph/extract/object/row/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ - -from . extract import * - diff --git a/trustgraph-flow/trustgraph/extract/object/row/extract.py b/trustgraph-flow/trustgraph/extract/object/row/extract.py deleted file mode 100755 index e262c1cb..00000000 --- a/trustgraph-flow/trustgraph/extract/object/row/extract.py +++ /dev/null @@ -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__) -