New nlp-query service

trustgraph-base/trustgraph/schema/services/nlp_query.py

@@ -7,16 +7,15 @@ from ..core.topic import topic
 
 # NLP to Structured Query Service - converts natural language to GraphQL
 
-class NLPToStructuredQueryRequest(Record):
+class QuestionToStructuredQueryRequest(Record):
-    natural_language_query = String()
+    question = String()
     max_results = Integer()
-    context_hints = Map(String())  # Optional context for query generation
 
-class NLPToStructuredQueryResponse(Record):
+class QuestionToStructuredQueryResponse(Record):
     error = Error()
     graphql_query = String()  # Generated GraphQL query
     variables = Map(String())  # GraphQL variables if any
     detected_schemas = Array(String())  # Which schemas the query targets
     confidence = Double()
 
-############################################################################
+############################################################################

trustgraph-flow/trustgraph/retrieval/nlp_query/__init__.py

@@ -0,0 +1 @@
+from . service import *

trustgraph-flow/trustgraph/retrieval/nlp_query/__main__.py

@@ -0,0 +1,5 @@
+#!/usr/bin/env python3
+
+from . service import run
+
+run()

trustgraph-flow/trustgraph/retrieval/nlp_query/service.py

@@ -0,0 +1,303 @@
+"""
+NLP to Structured Query Service - converts natural language questions to GraphQL queries.
+Two-phase approach: 1) Select relevant schemas, 2) Generate GraphQL query.
+"""
+
+import json
+import logging
+from typing import Dict, Any, Optional, List
+
+from ...schema import QuestionToStructuredQueryRequest, QuestionToStructuredQueryResponse
+from ...schema import PromptRequest, PromptResponse
+from ...schema import Error, RowSchema, Field as SchemaField
+
+from ...base import FlowProcessor, ConsumerSpec, ProducerSpec, ClientSpec
+
+# Module logger
+logger = logging.getLogger(__name__)
+
+default_ident = "nlp-query"
+default_schema_selection_template = "schema-selection"
+default_graphql_generation_template = "graphql-generation"
+
+class Processor(FlowProcessor):
+    
+    def __init__(self, **params):
+        
+        id = params.get("id", default_ident)
+        
+        # Config key for schemas
+        self.config_key = params.get("config_type", "schema")
+        
+        # Configurable prompt template names
+        self.schema_selection_template = params.get("schema_selection_template", default_schema_selection_template)
+        self.graphql_generation_template = params.get("graphql_generation_template", default_graphql_generation_template)
+        
+        super(Processor, self).__init__(
+            **params | {
+                "id": id,
+                "config-type": self.config_key,
+            }
+        )
+        
+        self.register_specification(
+            ConsumerSpec(
+                name = "request",
+                schema = QuestionToStructuredQueryRequest,
+                handler = self.on_message
+            )
+        )
+        
+        self.register_specification(
+            ProducerSpec(
+                name = "response",
+                schema = QuestionToStructuredQueryResponse,
+            )
+        )
+        
+        # Client spec for calling prompt service
+        self.register_specification(
+            ClientSpec(
+                request_name = "prompt-request",
+                response_name = "prompt-response",
+                request_schema = PromptRequest,
+                response_schema = PromptResponse
+            )
+        )
+        
+        # Register config handler for schema updates
+        self.register_config_handler(self.on_schema_config)
+        
+        # Schema storage: name -> RowSchema
+        self.schemas: Dict[str, RowSchema] = {}
+        
+        logger.info("NLP Query service initialized")
+
+    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 = {}
+        
+        # Check if our config type exists
+        if self.config_key not in config:
+            logger.warning(f"No '{self.config_key}' type in configuration")
+            return
+        
+        # Get the schemas dictionary for our type
+        schemas_config = config[self.config_key]
+        
+        # Process each schema in the schemas config
+        for schema_name, schema_json in schemas_config.items():
+            try:
+                # Parse the JSON schema definition
+                schema_def = json.loads(schema_json)
+                
+                # Create Field objects
+                fields = []
+                for field_def in schema_def.get("fields", []):
+                    field = SchemaField(
+                        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 phase1_select_schemas(self, question: str) -> List[str]:
+        """Phase 1: Use prompt service to select relevant schemas for the question"""
+        logger.info("Starting Phase 1: Schema selection")
+        
+        # Prepare schema information for the prompt
+        schema_info = []
+        for name, schema in self.schemas.items():
+            schema_desc = {
+                "name": name,
+                "description": schema.description,
+                "fields": [{"name": f.name, "type": f.type, "description": f.description} 
+                          for f in schema.fields]
+            }
+            schema_info.append(schema_desc)
+        
+        # Create prompt variables
+        variables = {
+            "question": question,
+            "schemas": json.dumps(schema_info, indent=2)
+        }
+        
+        # Call prompt service for schema selection
+        prompt_request = PromptRequest(
+            template=self.schema_selection_template,
+            variables=variables
+        )
+        
+        try:
+            response = await self.client("prompt-request").request(prompt_request)
+            
+            if response.error is not None:
+                raise Exception(f"Prompt service error: {response.error}")
+            
+            # Parse the response to get selected schema names
+            # Expecting response.text to contain JSON array of schema names
+            selected_schemas = json.loads(response.text)
+            
+            logger.info(f"Phase 1 selected schemas: {selected_schemas}")
+            return selected_schemas
+            
+        except Exception as e:
+            logger.error(f"Phase 1 schema selection failed: {e}")
+            raise
+
+    async def phase2_generate_graphql(self, question: str, selected_schemas: List[str]) -> Dict[str, Any]:
+        """Phase 2: Generate GraphQL query using selected schemas"""
+        logger.info(f"Starting Phase 2: GraphQL generation for schemas: {selected_schemas}")
+        
+        # Get detailed schema information for selected schemas only
+        selected_schema_info = []
+        for schema_name in selected_schemas:
+            if schema_name in self.schemas:
+                schema = self.schemas[schema_name]
+                schema_desc = {
+                    "name": schema_name,
+                    "description": schema.description,
+                    "fields": [
+                        {
+                            "name": f.name, 
+                            "type": f.type, 
+                            "description": f.description,
+                            "required": f.required,
+                            "primary_key": f.primary,
+                            "indexed": f.indexed,
+                            "enum_values": f.enum_values if f.enum_values else []
+                        } 
+                        for f in schema.fields
+                    ]
+                }
+                selected_schema_info.append(schema_desc)
+        
+        # Create prompt variables for GraphQL generation
+        variables = {
+            "question": question,
+            "schemas": json.dumps(selected_schema_info, indent=2)
+        }
+        
+        # Call prompt service for GraphQL generation
+        prompt_request = PromptRequest(
+            template=self.graphql_generation_template,
+            variables=variables
+        )
+        
+        try:
+            response = await self.client("prompt-request").request(prompt_request)
+            
+            if response.error is not None:
+                raise Exception(f"Prompt service error: {response.error}")
+            
+            # Parse the response to get GraphQL query and variables
+            # Expecting response.text to contain JSON with "query" and "variables" fields
+            result = json.loads(response.text)
+            
+            logger.info(f"Phase 2 generated GraphQL: {result.get('query', '')[:100]}...")
+            return result
+            
+        except Exception as e:
+            logger.error(f"Phase 2 GraphQL generation failed: {e}")
+            raise
+
+    async def on_message(self, msg, consumer, flow):
+        """Handle incoming question to structured query request"""
+        
+        try:
+            request = msg.value()
+            
+            # Sender-produced ID
+            id = msg.properties()["id"]
+            
+            logger.info(f"Handling NLP query request {id}: {request.question[:100]}...")
+            
+            # Phase 1: Select relevant schemas
+            selected_schemas = await self.phase1_select_schemas(request.question)
+            
+            # Phase 2: Generate GraphQL query
+            graphql_result = await self.phase2_generate_graphql(request.question, selected_schemas)
+            
+            # Create response
+            response = QuestionToStructuredQueryResponse(
+                error=None,
+                graphql_query=graphql_result.get("query", ""),
+                variables=graphql_result.get("variables", {}),
+                detected_schemas=selected_schemas,
+                confidence=graphql_result.get("confidence", 0.8)  # Default confidence
+            )
+            
+            logger.info("Sending NLP query response...")
+            await flow("response").send(response, properties={"id": id})
+            
+            logger.info("NLP query request completed")
+            
+        except Exception as e:
+            
+            logger.error(f"Exception in NLP query service: {e}", exc_info=True)
+            
+            logger.info("Sending error response...")
+            
+            response = QuestionToStructuredQueryResponse(
+                error = Error(
+                    type = "nlp-query-error",
+                    message = str(e),
+                ),
+                graphql_query = "",
+                variables = {},
+                detected_schemas = [],
+                confidence = 0.0
+            )
+            
+            await flow("response").send(response, properties={"id": id})
+
+    @staticmethod
+    def add_args(parser):
+        """Add command-line arguments"""
+        
+        FlowProcessor.add_args(parser)
+        
+        parser.add_argument(
+            '--config-type',
+            default='schema',
+            help='Configuration type prefix for schemas (default: schema)'
+        )
+        
+        parser.add_argument(
+            '--schema-selection-template',
+            default=default_schema_selection_template,
+            help=f'Prompt template name for schema selection (default: {default_schema_selection_template})'
+        )
+        
+        parser.add_argument(
+            '--graphql-generation-template', 
+            default=default_graphql_generation_template,
+            help=f'Prompt template name for GraphQL generation (default: {default_graphql_generation_template})'
+        )
+
+def run():
+    """Entry point for nlp-query command"""
+    Processor.launch(default_ident, __doc__)