trustgraph-base 1.4.19__tar.gz → 1.6.4__tar.gz

{trustgraph_base-1.4.19 → trustgraph_base-1.6.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-base
-Version: 1.4.19
+Version: 1.6.4
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Author-email: "trustgraph.ai" <security@trustgraph.ai>
 Project-URL: Homepage, https://github.com/trustgraph-ai/trustgraph

{trustgraph_base-1.4.19 → trustgraph_base-1.6.4}/trustgraph/base/__init__.py

@@ -12,7 +12,7 @@ from . parameter_spec import ParameterSpec
 from . producer_spec import ProducerSpec
 from . subscriber_spec import SubscriberSpec
 from . request_response_spec import RequestResponseSpec
-from . llm_service import LlmService, LlmResult
+from . llm_service import LlmService, LlmResult, LlmChunk
 from . chunking_service import ChunkingService
 from . embeddings_service import EmbeddingsService
 from . embeddings_client import EmbeddingsClientSpec

{trustgraph_base-1.4.19 → trustgraph_base-1.6.4}/trustgraph/base/embeddings_service.py

@@ -9,7 +9,7 @@ from prometheus_client import Histogram
 
 from .. schema import EmbeddingsRequest, EmbeddingsResponse, Error
 from .. exceptions import TooManyRequests
-from .. base import FlowProcessor, ConsumerSpec, ProducerSpec
+from .. base import FlowProcessor, ConsumerSpec, ProducerSpec, ParameterSpec
 
 # Module logger
 logger = logging.getLogger(__name__)
@@ -45,6 +45,12 @@ class EmbeddingsService(FlowProcessor):
             )
         )
 
+        self.register_specification(
+            ParameterSpec(
+                name = "model",
+            )
+        )
+
     async def on_request(self, msg, consumer, flow):
 
         try:
@@ -57,7 +63,9 @@ class EmbeddingsService(FlowProcessor):
 
             logger.debug(f"Handling embeddings request {id}...")
 
-            vectors = await self.on_embeddings(request.text)
+            # Pass model from request if specified (non-empty), otherwise use default
+            model = flow("model")
+            vectors = await self.on_embeddings(request.text, model=model)
 
             await flow("response").send(
                 EmbeddingsResponse(

{trustgraph_base-1.4.19 → trustgraph_base-1.6.4}/trustgraph/base/llm_service.py

@@ -28,6 +28,19 @@ class LlmResult:
         self.model = model
     __slots__ = ["text", "in_token", "out_token", "model"]
 
+class LlmChunk:
+    """Represents a streaming chunk from an LLM"""
+    def __init__(
+            self, text = None, in_token = None, out_token = None,
+            model = None, is_final = False,
+    ):
+        self.text = text
+        self.in_token = in_token
+        self.out_token = out_token
+        self.model = model
+        self.is_final = is_final
+    __slots__ = ["text", "in_token", "out_token", "model", "is_final"]
+
 class LlmService(FlowProcessor):
 
     def __init__(self, **params):
@@ -99,16 +112,57 @@ class LlmService(FlowProcessor):
 
             id = msg.properties()["id"]
 
-            with __class__.text_completion_metric.labels(
-                    id=self.id,
-                    flow=f"{flow.name}-{consumer.name}",
-            ).time():
-
-                model = flow("model")
-                temperature = flow("temperature")
-
-                response = await self.generate_content(
-                    request.system, request.prompt, model, temperature
+            model = flow("model")
+            temperature = flow("temperature")
+
+            # Check if streaming is requested and supported
+            streaming = getattr(request, 'streaming', False)
+
+            if streaming and self.supports_streaming():
+
+                # Streaming mode
+                with __class__.text_completion_metric.labels(
+                        id=self.id,
+                        flow=f"{flow.name}-{consumer.name}",
+                ).time():
+
+                    async for chunk in self.generate_content_stream(
+                        request.system, request.prompt, model, temperature
+                    ):
+                        await flow("response").send(
+                            TextCompletionResponse(
+                                error=None,
+                                response=chunk.text,
+                                in_token=chunk.in_token,
+                                out_token=chunk.out_token,
+                                model=chunk.model,
+                                end_of_stream=chunk.is_final
+                            ),
+                            properties={"id": id}
+                        )
+
+            else:
+
+                # Non-streaming mode (original behavior)
+                with __class__.text_completion_metric.labels(
+                        id=self.id,
+                        flow=f"{flow.name}-{consumer.name}",
+                ).time():
+
+                    response = await self.generate_content(
+                        request.system, request.prompt, model, temperature
+                    )
+
+                await flow("response").send(
+                    TextCompletionResponse(
+                        error=None,
+                        response=response.text,
+                        in_token=response.in_token,
+                        out_token=response.out_token,
+                        model=response.model,
+                        end_of_stream=True
+                    ),
+                    properties={"id": id}
                 )
 
             __class__.text_completion_model_metric.labels(
@@ -119,17 +173,6 @@ class LlmService(FlowProcessor):
                 "temperature": str(temperature) if temperature is not None else "",
             })
 
-            await flow("response").send(
-                TextCompletionResponse(
-                    error=None,
-                    response=response.text,
-                    in_token=response.in_token,
-                    out_token=response.out_token,
-                    model=response.model
-                ),
-                properties={"id": id}
-            )
-
         except TooManyRequests as e:
             raise e
 
@@ -151,10 +194,26 @@ class LlmService(FlowProcessor):
                     in_token=None,
                     out_token=None,
                     model=None,
+                    end_of_stream=True
                 ),
                 properties={"id": id}
             )
 
+    def supports_streaming(self):
+        """
+        Override in subclass to indicate streaming support.
+        Returns False by default.
+        """
+        return False
+
+    async def generate_content_stream(self, system, prompt, model=None, temperature=None):
+        """
+        Override in subclass to implement streaming.
+        Should yield LlmChunk objects.
+        The final chunk should have is_final=True.
+        """
+        raise NotImplementedError("Streaming not implemented for this provider")
+
     @staticmethod
     def add_args(parser):
 

trustgraph_base-1.6.4/trustgraph/base/prompt_client.py

@@ -0,0 +1,171 @@
+
+import json
+import asyncio
+import logging
+
+from . request_response_spec import RequestResponse, RequestResponseSpec
+from .. schema import PromptRequest, PromptResponse
+
+logger = logging.getLogger(__name__)
+
+class PromptClient(RequestResponse):
+
+    async def prompt(self, id, variables, timeout=600, streaming=False, chunk_callback=None):
+        logger.info(f"DEBUG prompt_client: prompt called, id={id}, streaming={streaming}, chunk_callback={chunk_callback is not None}")
+
+        if not streaming:
+            logger.info("DEBUG prompt_client: Non-streaming path")
+            # Non-streaming path
+            resp = await self.request(
+                PromptRequest(
+                    id = id,
+                    terms = {
+                        k: json.dumps(v)
+                        for k, v in variables.items()
+                    },
+                    streaming = False
+                ),
+                timeout=timeout
+            )
+
+            if resp.error:
+                raise RuntimeError(resp.error.message)
+
+            if resp.text: return resp.text
+
+            return json.loads(resp.object)
+
+        else:
+            logger.info("DEBUG prompt_client: Streaming path")
+            # Streaming path - collect all chunks
+            full_text = ""
+            full_object = None
+
+            async def collect_chunks(resp):
+                nonlocal full_text, full_object
+                logger.info(f"DEBUG prompt_client: collect_chunks called, resp.text={resp.text[:50] if resp.text else None}, end_of_stream={getattr(resp, 'end_of_stream', False)}")
+
+                if resp.error:
+                    logger.error(f"DEBUG prompt_client: Error in response: {resp.error.message}")
+                    raise RuntimeError(resp.error.message)
+
+                if resp.text:
+                    full_text += resp.text
+                    logger.info(f"DEBUG prompt_client: Accumulated {len(full_text)} chars")
+                    # Call chunk callback if provided
+                    if chunk_callback:
+                        logger.info(f"DEBUG prompt_client: Calling chunk_callback")
+                        if asyncio.iscoroutinefunction(chunk_callback):
+                            await chunk_callback(resp.text)
+                        else:
+                            chunk_callback(resp.text)
+                elif resp.object:
+                    logger.info(f"DEBUG prompt_client: Got object response")
+                    full_object = resp.object
+
+                end_stream = getattr(resp, 'end_of_stream', False)
+                logger.info(f"DEBUG prompt_client: Returning end_of_stream={end_stream}")
+                return end_stream
+
+            logger.info("DEBUG prompt_client: Creating PromptRequest")
+            req = PromptRequest(
+                id = id,
+                terms = {
+                    k: json.dumps(v)
+                    for k, v in variables.items()
+                },
+                streaming = True
+            )
+            logger.info(f"DEBUG prompt_client: About to call self.request with recipient, timeout={timeout}")
+            await self.request(
+                req,
+                recipient=collect_chunks,
+                timeout=timeout
+            )
+            logger.info(f"DEBUG prompt_client: self.request returned, full_text has {len(full_text)} chars")
+
+            if full_text:
+                logger.info("DEBUG prompt_client: Returning full_text")
+                return full_text
+
+            logger.info("DEBUG prompt_client: Returning parsed full_object")
+            return json.loads(full_object)
+
+    async def extract_definitions(self, text, timeout=600):
+        return await self.prompt(
+            id = "extract-definitions",
+            variables = { "text": text },
+            timeout = timeout,
+        )
+
+    async def extract_relationships(self, text, timeout=600):
+        return await self.prompt(
+            id = "extract-relationships",
+            variables = { "text": text },
+            timeout = timeout,
+        )
+
+    async def extract_objects(self, text, schema, timeout=600):
+        return await self.prompt(
+            id = "extract-rows",
+            variables = { "text": text, "schema": schema, },
+            timeout = timeout,
+        )
+
+    async def kg_prompt(self, query, kg, timeout=600, streaming=False, chunk_callback=None):
+        return await self.prompt(
+            id = "kg-prompt",
+            variables = {
+                "query": query,
+                "knowledge": [
+                    { "s": v[0], "p": v[1], "o": v[2] }
+                    for v in kg
+                ]
+            },
+            timeout = timeout,
+            streaming = streaming,
+            chunk_callback = chunk_callback,
+        )
+
+    async def document_prompt(self, query, documents, timeout=600, streaming=False, chunk_callback=None):
+        return await self.prompt(
+            id = "document-prompt",
+            variables = {
+                "query": query,
+                "documents": documents,
+            },
+            timeout = timeout,
+            streaming = streaming,
+            chunk_callback = chunk_callback,
+        )
+
+    async def agent_react(self, variables, timeout=600, streaming=False, chunk_callback=None):
+        return await self.prompt(
+            id = "agent-react",
+            variables = variables,
+            timeout = timeout,
+            streaming = streaming,
+            chunk_callback = chunk_callback,
+        )
+
+    async def question(self, question, timeout=600):
+        return await self.prompt(
+            id = "question",
+            variables = {
+                "question": question,
+            },
+            timeout = timeout,
+        )
+
+class PromptClientSpec(RequestResponseSpec):
+    def __init__(
+            self, request_name, response_name,
+    ):
+        super(PromptClientSpec, self).__init__(
+            request_name = request_name,
+            request_schema = PromptRequest,
+            response_name = response_name,
+            response_schema = PromptResponse,
+            impl = PromptClient,
+        )
+

{trustgraph_base-1.4.19 → trustgraph_base-1.6.4}/trustgraph/base/subscriber.py

@@ -43,12 +43,18 @@ class Subscriber:
 
     async def start(self):
 
-        self.consumer = self.client.subscribe(
-            topic = self.topic,
-            subscription_name = self.subscription,
-            consumer_name = self.consumer_name,
-            schema = JsonSchema(self.schema),
-        )
+        # Build subscribe arguments
+        subscribe_args = {
+            'topic': self.topic,
+            'subscription_name': self.subscription,
+            'consumer_name': self.consumer_name,
+        }
+
+        # Only add schema if provided (omit if None)
+        if self.schema is not None:
+            subscribe_args['schema'] = JsonSchema(self.schema)
+
+        self.consumer = self.client.subscribe(**subscribe_args)
 
         self.task = asyncio.create_task(self.run())
 
@@ -87,10 +93,14 @@ class Subscriber:
                     if self.draining and drain_end_time is None:
                         drain_end_time = time.time() + self.drain_timeout
                         logger.info(f"Subscriber entering drain mode, timeout={self.drain_timeout}s")
-                        
+
                         # Stop accepting new messages from Pulsar during drain
                         if self.consumer:
-                            self.consumer.pause_message_listener()
+                            try:
+                                self.consumer.pause_message_listener()
+                            except _pulsar.InvalidConfiguration:
+                                # Not all consumers have message listeners (e.g., blocking receive mode)
+                                pass
                     
                     # Check drain timeout
                     if self.draining and drain_end_time and time.time() > drain_end_time:
@@ -145,12 +155,21 @@ class Subscriber:
             finally:
                 # Negative acknowledge any pending messages
                 for msg in self.pending_acks.values():
-                    self.consumer.negative_acknowledge(msg)
+                    try:
+                        self.consumer.negative_acknowledge(msg)
+                    except _pulsar.AlreadyClosed:
+                        pass  # Consumer already closed
                 self.pending_acks.clear()
 
                 if self.consumer:
-                    self.consumer.unsubscribe()
-                    self.consumer.close()
+                    try:
+                        self.consumer.unsubscribe()
+                    except _pulsar.AlreadyClosed:
+                        pass  # Already closed
+                    try:
+                        self.consumer.close()
+                    except _pulsar.AlreadyClosed:
+                        pass  # Already closed
                     self.consumer = None
                 
          

trustgraph_base-1.6.4/trustgraph/base/text_completion_client.py

@@ -0,0 +1,57 @@
+
+from . request_response_spec import RequestResponse, RequestResponseSpec
+from .. schema import TextCompletionRequest, TextCompletionResponse
+
+class TextCompletionClient(RequestResponse):
+    async def text_completion(self, system, prompt, streaming=False, timeout=600):
+        # If not streaming, use original behavior
+        if not streaming:
+            resp = await self.request(
+                TextCompletionRequest(
+                    system = system, prompt = prompt, streaming = False
+                ),
+                timeout=timeout
+            )
+
+            if resp.error:
+                raise RuntimeError(resp.error.message)
+
+            return resp.response
+
+        # For streaming: collect all chunks and return complete response
+        full_response = ""
+
+        async def collect_chunks(resp):
+            nonlocal full_response
+
+            if resp.error:
+                raise RuntimeError(resp.error.message)
+
+            if resp.response:
+                full_response += resp.response
+
+            # Return True when end_of_stream is reached
+            return getattr(resp, 'end_of_stream', False)
+
+        await self.request(
+            TextCompletionRequest(
+                system = system, prompt = prompt, streaming = True
+            ),
+            recipient=collect_chunks,
+            timeout=timeout
+        )
+
+        return full_response
+
+class TextCompletionClientSpec(RequestResponseSpec):
+    def __init__(
+            self, request_name, response_name,
+    ):
+        super(TextCompletionClientSpec, self).__init__(
+            request_name = request_name,
+            request_schema = TextCompletionRequest,
+            response_name = response_name,
+            response_schema = TextCompletionResponse,
+            impl = TextCompletionClient,
+        )
+

trustgraph_base-1.6.4/trustgraph/base_version.py

@@ -0,0 +1 @@
+__version__ = "1.6.4"

trustgraph_base-1.6.4/trustgraph/clients/llm_client.py

@@ -0,0 +1,105 @@
+
+import _pulsar
+
+from .. schema import TextCompletionRequest, TextCompletionResponse
+from .. schema import text_completion_request_queue
+from .. schema import text_completion_response_queue
+from . base import BaseClient
+from .. exceptions import LlmError
+
+# Ugly
+ERROR=_pulsar.LoggerLevel.Error
+WARN=_pulsar.LoggerLevel.Warn
+INFO=_pulsar.LoggerLevel.Info
+DEBUG=_pulsar.LoggerLevel.Debug
+
+class LlmClient(BaseClient):
+
+    def __init__(
+            self, log_level=ERROR,
+            subscriber=None,
+            input_queue=None,
+            output_queue=None,
+            pulsar_host="pulsar://pulsar:6650",
+            pulsar_api_key=None,
+    ):
+
+        if input_queue is None: input_queue = text_completion_request_queue
+        if output_queue is None: output_queue = text_completion_response_queue
+
+        super(LlmClient, self).__init__(
+            log_level=log_level,
+            subscriber=subscriber,
+            input_queue=input_queue,
+            output_queue=output_queue,
+            pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
+            input_schema=TextCompletionRequest,
+            output_schema=TextCompletionResponse,
+        )
+
+    def request(self, system, prompt, timeout=300, streaming=False):
+        """
+        Non-streaming request (backward compatible).
+        Returns complete response string.
+        """
+        if streaming:
+            raise ValueError("Use request_stream() for streaming requests")
+        return self.call(
+            system=system, prompt=prompt, streaming=False, timeout=timeout
+        ).response
+
+    def request_stream(self, system, prompt, timeout=300):
+        """
+        Streaming request generator.
+        Yields response chunks as they arrive.
+        Usage:
+            for chunk in client.request_stream(system, prompt):
+                print(chunk.response, end='', flush=True)
+        """
+        import time
+        import uuid
+
+        id = str(uuid.uuid4())
+        request = TextCompletionRequest(
+            system=system, prompt=prompt, streaming=True
+        )
+
+        end_time = time.time() + timeout
+        self.producer.send(request, properties={"id": id})
+
+        # Collect responses until end_of_stream
+        while time.time() < end_time:
+            try:
+                msg = self.consumer.receive(timeout_millis=2500)
+            except Exception:
+                continue
+
+            mid = msg.properties()["id"]
+
+            if mid == id:
+                value = msg.value()
+
+                # Handle errors
+                if value.error:
+                    self.consumer.acknowledge(msg)
+                    if value.error.type == "llm-error":
+                        raise LlmError(value.error.message)
+                    else:
+                        raise RuntimeError(
+                            f"{value.error.type}: {value.error.message}"
+                        )
+
+                self.consumer.acknowledge(msg)
+                yield value
+
+                # Check if this is the final chunk
+                if getattr(value, 'end_of_stream', True):
+                    break
+            else:
+                # Ignore messages with wrong ID
+                self.consumer.acknowledge(msg)
+
+        if time.time() >= end_time:
+            raise TimeoutError("Timed out waiting for response")
+

trustgraph_base-1.6.4/trustgraph/messaging/translators/agent.py

@@ -0,0 +1,70 @@
+from typing import Dict, Any, Tuple
+from ...schema import AgentRequest, AgentResponse
+from .base import MessageTranslator
+
+
+class AgentRequestTranslator(MessageTranslator):
+    """Translator for AgentRequest schema objects"""
+    
+    def to_pulsar(self, data: Dict[str, Any]) -> AgentRequest:
+        return AgentRequest(
+            question=data["question"],
+            state=data.get("state", None),
+            group=data.get("group", None),
+            history=data.get("history", []),
+            user=data.get("user", "trustgraph"),
+            streaming=data.get("streaming", False)
+        )
+
+    def from_pulsar(self, obj: AgentRequest) -> Dict[str, Any]:
+        return {
+            "question": obj.question,
+            "state": obj.state,
+            "group": obj.group,
+            "history": obj.history,
+            "user": obj.user,
+            "streaming": getattr(obj, "streaming", False)
+        }
+
+
+class AgentResponseTranslator(MessageTranslator):
+    """Translator for AgentResponse schema objects"""
+    
+    def to_pulsar(self, data: Dict[str, Any]) -> AgentResponse:
+        raise NotImplementedError("Response translation to Pulsar not typically needed")
+    
+    def from_pulsar(self, obj: AgentResponse) -> Dict[str, Any]:
+        result = {}
+
+        # Check if this is a streaming response (has chunk_type)
+        if hasattr(obj, 'chunk_type') and obj.chunk_type:
+            result["chunk_type"] = obj.chunk_type
+            if obj.content:
+                result["content"] = obj.content
+            result["end_of_message"] = getattr(obj, "end_of_message", False)
+            result["end_of_dialog"] = getattr(obj, "end_of_dialog", False)
+        else:
+            # Legacy format
+            if obj.answer:
+                result["answer"] = obj.answer
+            if obj.thought:
+                result["thought"] = obj.thought
+            if obj.observation:
+                result["observation"] = obj.observation
+
+        # Always include error if present
+        if hasattr(obj, 'error') and obj.error and obj.error.message:
+            result["error"] = {"message": obj.error.message, "code": obj.error.code}
+
+        return result
+
+    def from_response_with_completion(self, obj: AgentResponse) -> Tuple[Dict[str, Any], bool]:
+        """Returns (response_dict, is_final)"""
+        # For streaming responses, check end_of_dialog
+        if hasattr(obj, 'chunk_type') and obj.chunk_type:
+            is_final = getattr(obj, 'end_of_dialog', False)
+        else:
+            # For legacy responses, check if answer is present
+            is_final = (obj.answer is not None)
+
+        return self.from_pulsar(obj), is_final

{trustgraph_base-1.4.19 → trustgraph_base-1.6.4}/trustgraph/messaging/translators/prompt.py

@@ -16,10 +16,11 @@ class PromptRequestTranslator(MessageTranslator):
                 k: json.dumps(v)
                 for k, v in data["variables"].items()
             }
-        
+
         return PromptRequest(
             id=data.get("id"),
-            terms=terms
+            terms=terms,
+            streaming=data.get("streaming", False)
         )
     
     def from_pulsar(self, obj: PromptRequest) -> Dict[str, Any]:
@@ -51,4 +52,6 @@ class PromptResponseTranslator(MessageTranslator):
     
     def from_response_with_completion(self, obj: PromptResponse) -> Tuple[Dict[str, Any], bool]:
         """Returns (response_dict, is_final)"""
-        return self.from_pulsar(obj), True
+        # Check end_of_stream field to determine if this is the final message
+        is_final = getattr(obj, 'end_of_stream', True)
+        return self.from_pulsar(obj), is_final