mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 00:16:23 +02:00
Feature/streaming llm phase 1 (#566)
* Tidy up duplicate tech specs in doc directory * Streaming LLM text-completion service tech spec. * text-completion and prompt interfaces * streaming change applied to all LLMs, so far tested with VertexAI * Skip Pinecone unit tests, upstream module issue is affecting things, tests are passing again * Added agent streaming, not working and has broken tests
This commit is contained in:
cybermaggedon
GitHub
parent
943a9d83b0
commit
310a2deb06
44 changed files with 2684 additions and 937 deletions
|
|
@ -1,106 +0,0 @@
|
|||
# Knowledge Graph Architecture Foundations
|
||||
|
||||
## Foundation 1: Subject-Predicate-Object (SPO) Graph Model
|
||||
**Decision**: Adopt SPO/RDF as the core knowledge representation model
|
||||
|
||||
**Rationale**:
|
||||
- Provides maximum flexibility and interoperability with existing graph technologies
|
||||
- Enables seamless translation to other graph query languages (e.g., SPO → Cypher, but not vice versa)
|
||||
- Creates a foundation that "unlocks a lot" of downstream capabilities
|
||||
- Supports both node-to-node relationships (SPO) and node-to-literal relationships (RDF)
|
||||
|
||||
**Implementation**:
|
||||
- Core data structure: `node → edge → {node | literal}`
|
||||
- Maintain compatibility with RDF standards while supporting extended SPO operations
|
||||
|
||||
## Foundation 2: LLM-Native Knowledge Graph Integration
|
||||
**Decision**: Optimize knowledge graph structure and operations for LLM interaction
|
||||
|
||||
**Rationale**:
|
||||
- Primary use case involves LLMs interfacing with knowledge graphs
|
||||
- Graph technology choices must prioritize LLM compatibility over other considerations
|
||||
- Enables natural language processing workflows that leverage structured knowledge
|
||||
|
||||
**Implementation**:
|
||||
- Design graph schemas that LLMs can effectively reason about
|
||||
- Optimize for common LLM interaction patterns
|
||||
|
||||
## Foundation 3: Embedding-Based Graph Navigation
|
||||
**Decision**: Implement direct mapping from natural language queries to graph nodes via embeddings
|
||||
|
||||
**Rationale**:
|
||||
- Enables the simplest possible path from NLP query to graph navigation
|
||||
- Avoids complex intermediate query generation steps
|
||||
- Provides efficient semantic search capabilities within the graph structure
|
||||
|
||||
**Implementation**:
|
||||
- `NLP Query → Graph Embeddings → Graph Nodes`
|
||||
- Maintain embedding representations for all graph entities
|
||||
- Support direct semantic similarity matching for query resolution
|
||||
|
||||
## Foundation 4: Distributed Entity Resolution with Deterministic Identifiers
|
||||
**Decision**: Support parallel knowledge extraction with deterministic entity identification (80% rule)
|
||||
|
||||
**Rationale**:
|
||||
- **Ideal**: Single-process extraction with complete state visibility enables perfect entity resolution
|
||||
- **Reality**: Scalability requirements demand parallel processing capabilities
|
||||
- **Compromise**: Design for deterministic entity identification across distributed processes
|
||||
|
||||
**Implementation**:
|
||||
- Develop mechanisms for generating consistent, unique identifiers across different knowledge extractors
|
||||
- Same entity mentioned in different processes must resolve to the same identifier
|
||||
- Acknowledge that ~20% of edge cases may require alternative processing models
|
||||
- Design fallback mechanisms for complex entity resolution scenarios
|
||||
|
||||
## Foundation 5: Event-Driven Architecture with Publish-Subscribe
|
||||
**Decision**: Implement pub-sub messaging system for system coordination
|
||||
|
||||
**Rationale**:
|
||||
- Enables loose coupling between knowledge extraction, storage, and query components
|
||||
- Supports real-time updates and notifications across the system
|
||||
- Facilitates scalable, distributed processing workflows
|
||||
|
||||
**Implementation**:
|
||||
- Message-driven coordination between system components
|
||||
- Event streams for knowledge updates, extraction completion, and query results
|
||||
|
||||
## Foundation 6: Reentrant Agent Communication
|
||||
**Decision**: Support reentrant pub-sub operations for agent-based processing
|
||||
|
||||
**Rationale**:
|
||||
- Enables sophisticated agent workflows where agents can trigger and respond to each other
|
||||
- Supports complex, multi-step knowledge processing pipelines
|
||||
- Allows for recursive and iterative processing patterns
|
||||
|
||||
**Implementation**:
|
||||
- Pub-sub system must handle reentrant calls safely
|
||||
- Agent coordination mechanisms that prevent infinite loops
|
||||
- Support for agent workflow orchestration
|
||||
|
||||
## Foundation 7: Columnar Data Store Integration
|
||||
**Decision**: Ensure query compatibility with columnar storage systems
|
||||
|
||||
**Rationale**:
|
||||
- Enables efficient analytical queries over large knowledge datasets
|
||||
- Supports business intelligence and reporting use cases
|
||||
- Bridges graph-based knowledge representation with traditional analytical workflows
|
||||
|
||||
**Implementation**:
|
||||
- Query translation layer: Graph queries → Columnar queries
|
||||
- Hybrid storage strategy supporting both graph operations and analytical workloads
|
||||
- Maintain query performance across both paradigms
|
||||
|
||||
---
|
||||
|
||||
## Architecture Principles Summary
|
||||
|
||||
1. **Flexibility First**: SPO/RDF model provides maximum adaptability
|
||||
2. **LLM Optimization**: All design decisions consider LLM interaction requirements
|
||||
3. **Semantic Efficiency**: Direct embedding-to-node mapping for optimal query performance
|
||||
4. **Pragmatic Scalability**: Balance perfect accuracy with practical distributed processing
|
||||
5. **Event-Driven Coordination**: Pub-sub enables loose coupling and scalability
|
||||
6. **Agent-Friendly**: Support complex, multi-agent processing workflows
|
||||
7. **Analytical Compatibility**: Bridge graph and columnar paradigms for comprehensive querying
|
||||
|
||||
These foundations establish a knowledge graph architecture that balances theoretical rigor with practical scalability requirements, optimized for LLM integration and distributed processing.
|
||||
|
||||
|
|
@ -1,169 +0,0 @@
|
|||
# TrustGraph Logging Strategy
|
||||
|
||||
## Overview
|
||||
|
||||
TrustGraph uses Python's built-in `logging` module for all logging operations. This provides a standardized, flexible approach to logging across all components of the system.
|
||||
|
||||
## Default Configuration
|
||||
|
||||
### Logging Level
|
||||
- **Default Level**: `INFO`
|
||||
- **Debug Mode**: `DEBUG` (enabled via command-line argument)
|
||||
- **Production**: `WARNING` or `ERROR` as appropriate
|
||||
|
||||
### Output Destination
|
||||
All logs should be written to **standard output (stdout)** to ensure compatibility with containerized environments and log aggregation systems.
|
||||
|
||||
## Implementation Guidelines
|
||||
|
||||
### 1. Logger Initialization
|
||||
|
||||
Each module should create its own logger using the module's `__name__`:
|
||||
|
||||
```python
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
```
|
||||
|
||||
### 2. Centralized Configuration
|
||||
|
||||
The logging configuration should be centralized in `async_processor.py` (or a dedicated logging configuration module) since it's inherited by much of the codebase:
|
||||
|
||||
```python
|
||||
import logging
|
||||
import argparse
|
||||
|
||||
def setup_logging(log_level='INFO'):
|
||||
"""Configure logging for the entire application"""
|
||||
logging.basicConfig(
|
||||
level=getattr(logging, log_level.upper()),
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
||||
handlers=[logging.StreamHandler()]
|
||||
)
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
'--log-level',
|
||||
default='INFO',
|
||||
choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],
|
||||
help='Set the logging level (default: INFO)'
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
# In main execution
|
||||
if __name__ == '__main__':
|
||||
args = parse_args()
|
||||
setup_logging(args.log_level)
|
||||
```
|
||||
|
||||
### 3. Logging Best Practices
|
||||
|
||||
#### Log Levels Usage
|
||||
- **DEBUG**: Detailed information for diagnosing problems (variable values, function entry/exit)
|
||||
- **INFO**: General informational messages (service started, configuration loaded, processing milestones)
|
||||
- **WARNING**: Warning messages for potentially harmful situations (deprecated features, recoverable errors)
|
||||
- **ERROR**: Error messages for serious problems (failed operations, exceptions)
|
||||
- **CRITICAL**: Critical messages for system failures requiring immediate attention
|
||||
|
||||
#### Message Format
|
||||
```python
|
||||
# Good - includes context
|
||||
logger.info(f"Processing document: {doc_id}, size: {doc_size} bytes")
|
||||
logger.error(f"Failed to connect to database: {error}", exc_info=True)
|
||||
|
||||
# Avoid - lacks context
|
||||
logger.info("Processing document")
|
||||
logger.error("Connection failed")
|
||||
```
|
||||
|
||||
#### Performance Considerations
|
||||
```python
|
||||
# Use lazy formatting for expensive operations
|
||||
logger.debug("Expensive operation result: %s", expensive_function())
|
||||
|
||||
# Check log level for very expensive debug operations
|
||||
if logger.isEnabledFor(logging.DEBUG):
|
||||
debug_data = compute_expensive_debug_info()
|
||||
logger.debug(f"Debug data: {debug_data}")
|
||||
```
|
||||
|
||||
### 4. Structured Logging
|
||||
|
||||
For complex data, use structured logging:
|
||||
|
||||
```python
|
||||
logger.info("Request processed", extra={
|
||||
'request_id': request_id,
|
||||
'duration_ms': duration,
|
||||
'status_code': status_code,
|
||||
'user_id': user_id
|
||||
})
|
||||
```
|
||||
|
||||
### 5. Exception Logging
|
||||
|
||||
Always include stack traces for exceptions:
|
||||
|
||||
```python
|
||||
try:
|
||||
process_data()
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to process data: {e}", exc_info=True)
|
||||
raise
|
||||
```
|
||||
|
||||
### 6. Async Logging Considerations
|
||||
|
||||
For async code, ensure thread-safe logging:
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
async def async_operation():
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info(f"Starting async operation in task: {asyncio.current_task().get_name()}")
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
|
||||
Support environment-based configuration as a fallback:
|
||||
|
||||
```python
|
||||
import os
|
||||
|
||||
log_level = os.environ.get('TRUSTGRAPH_LOG_LEVEL', 'INFO')
|
||||
```
|
||||
|
||||
## Testing
|
||||
|
||||
During tests, consider using a different logging configuration:
|
||||
|
||||
```python
|
||||
# In test setup
|
||||
logging.getLogger().setLevel(logging.WARNING) # Reduce noise during tests
|
||||
```
|
||||
|
||||
## Monitoring Integration
|
||||
|
||||
Ensure log format is compatible with monitoring tools:
|
||||
- Include timestamps in ISO format
|
||||
- Use consistent field names
|
||||
- Include correlation IDs where applicable
|
||||
- Structure logs for easy parsing (JSON format for production)
|
||||
|
||||
## Security Considerations
|
||||
|
||||
- Never log sensitive information (passwords, API keys, personal data)
|
||||
- Sanitize user input before logging
|
||||
- Use placeholders for sensitive fields: `user_id=****1234`
|
||||
|
||||
## Migration Path
|
||||
|
||||
For existing code using print statements:
|
||||
1. Replace `print()` with appropriate logger calls
|
||||
2. Choose appropriate log levels based on message importance
|
||||
3. Add context to make logs more useful
|
||||
4. Test logging output at different levels
|
||||
|
|
@ -1,91 +0,0 @@
|
|||
# Schema Directory Refactoring Proposal
|
||||
|
||||
## Current Issues
|
||||
|
||||
1. **Flat structure** - All schemas in one directory makes it hard to understand relationships
|
||||
2. **Mixed concerns** - Core types, domain objects, and API contracts all mixed together
|
||||
3. **Unclear naming** - Files like "object.py", "types.py", "topic.py" don't clearly indicate their purpose
|
||||
4. **No clear layering** - Can't easily see what depends on what
|
||||
|
||||
## Proposed Structure
|
||||
|
||||
```
|
||||
trustgraph-base/trustgraph/schema/
|
||||
├── __init__.py
|
||||
├── core/ # Core primitive types used everywhere
|
||||
│ ├── __init__.py
|
||||
│ ├── primitives.py # Error, Value, Triple, Field, RowSchema
|
||||
│ ├── metadata.py # Metadata record
|
||||
│ └── topic.py # Topic utilities
|
||||
│
|
||||
├── knowledge/ # Knowledge domain models and extraction
|
||||
│ ├── __init__.py
|
||||
│ ├── graph.py # EntityContext, EntityEmbeddings, Triples
|
||||
│ ├── document.py # Document, TextDocument, Chunk
|
||||
│ ├── knowledge.py # Knowledge extraction types
|
||||
│ ├── embeddings.py # All embedding-related types (moved from multiple files)
|
||||
│ └── nlp.py # Definition, Topic, Relationship, Fact types
|
||||
│
|
||||
└── services/ # Service request/response contracts
|
||||
├── __init__.py
|
||||
├── llm.py # TextCompletion, Embeddings, Tool requests/responses
|
||||
├── retrieval.py # GraphRAG, DocumentRAG queries/responses
|
||||
├── query.py # GraphEmbeddingsRequest/Response, DocumentEmbeddingsRequest/Response
|
||||
├── agent.py # Agent requests/responses
|
||||
├── flow.py # Flow requests/responses
|
||||
├── prompt.py # Prompt service requests/responses
|
||||
├── config.py # Configuration service
|
||||
├── library.py # Librarian service
|
||||
└── lookup.py # Lookup service
|
||||
```
|
||||
|
||||
## Key Changes
|
||||
|
||||
1. **Hierarchical organization** - Clear separation between core types, knowledge models, and service contracts
|
||||
2. **Better naming**:
|
||||
- `types.py` → `core/primitives.py` (clearer purpose)
|
||||
- `object.py` → Split between appropriate files based on actual content
|
||||
- `documents.py` → `knowledge/document.py` (singular, consistent)
|
||||
- `models.py` → `services/llm.py` (clearer what kind of models)
|
||||
- `prompt.py` → Split: service parts to `services/prompt.py`, data types to `knowledge/nlp.py`
|
||||
|
||||
3. **Logical grouping**:
|
||||
- All embedding types consolidated in `knowledge/embeddings.py`
|
||||
- All LLM-related service contracts in `services/llm.py`
|
||||
- Clear separation of request/response pairs in services directory
|
||||
- Knowledge extraction types grouped with other knowledge domain models
|
||||
|
||||
4. **Dependency clarity**:
|
||||
- Core types have no dependencies
|
||||
- Knowledge models depend only on core
|
||||
- Service contracts can depend on both core and knowledge models
|
||||
|
||||
## Migration Benefits
|
||||
|
||||
1. **Easier navigation** - Developers can quickly find what they need
|
||||
2. **Better modularity** - Clear boundaries between different concerns
|
||||
3. **Simpler imports** - More intuitive import paths
|
||||
4. **Future-proof** - Easy to add new knowledge types or services without cluttering
|
||||
|
||||
## Example Import Changes
|
||||
|
||||
```python
|
||||
# Before
|
||||
from trustgraph.schema import Error, Triple, GraphEmbeddings, TextCompletionRequest
|
||||
|
||||
# After
|
||||
from trustgraph.schema.core import Error, Triple
|
||||
from trustgraph.schema.knowledge import GraphEmbeddings
|
||||
from trustgraph.schema.services import TextCompletionRequest
|
||||
```
|
||||
|
||||
## Implementation Notes
|
||||
|
||||
1. Keep backward compatibility by maintaining imports in root `__init__.py`
|
||||
2. Move files gradually, updating imports as needed
|
||||
3. Consider adding a `legacy.py` that imports everything for transition period
|
||||
4. Update documentation to reflect new structure
|
||||
|
||||
<function_calls>
|
||||
<invoke name="TodoWrite">
|
||||
<parameter name="todos">[{"id": "1", "content": "Examine current schema directory structure", "status": "completed", "priority": "high"}, {"id": "2", "content": "Analyze schema files and their purposes", "status": "completed", "priority": "high"}, {"id": "3", "content": "Propose improved naming and structure", "status": "completed", "priority": "high"}]
|
||||
|
|
@ -1,253 +0,0 @@
|
|||
# Structured Data Technical Specification
|
||||
|
||||
## Overview
|
||||
|
||||
This specification describes the integration of TrustGraph with structured data flows, enabling the system to work with data that can be represented as rows in tables or objects in object stores. The integration supports four primary use cases:
|
||||
|
||||
1. **Unstructured to Structured Extraction**: Read unstructured data sources, identify and extract object structures, and store them in a tabular format
|
||||
2. **Structured Data Ingestion**: Load data that is already in structured formats directly into the structured store alongside extracted data
|
||||
3. **Natural Language Querying**: Convert natural language questions into structured queries to extract matching data from the store
|
||||
4. **Direct Structured Querying**: Execute structured queries directly against the data store for precise data retrieval
|
||||
|
||||
## Goals
|
||||
|
||||
- **Unified Data Access**: Provide a single interface for accessing both structured and unstructured data within TrustGraph
|
||||
- **Seamless Integration**: Enable smooth interoperability between TrustGraph's graph-based knowledge representation and traditional structured data formats
|
||||
- **Flexible Extraction**: Support automatic extraction of structured data from various unstructured sources (documents, text, etc.)
|
||||
- **Query Versatility**: Allow users to query data using both natural language and structured query languages
|
||||
- **Data Consistency**: Maintain data integrity and consistency across different data representations
|
||||
- **Performance Optimization**: Ensure efficient storage and retrieval of structured data at scale
|
||||
- **Schema Flexibility**: Support both schema-on-write and schema-on-read approaches to accommodate diverse data sources
|
||||
- **Backwards Compatibility**: Preserve existing TrustGraph functionality while adding structured data capabilities
|
||||
|
||||
## Background
|
||||
|
||||
TrustGraph currently excels at processing unstructured data and building knowledge graphs from diverse sources. However, many enterprise use cases involve data that is inherently structured - customer records, transaction logs, inventory databases, and other tabular datasets. These structured datasets often need to be analyzed alongside unstructured content to provide comprehensive insights.
|
||||
|
||||
Current limitations include:
|
||||
- No native support for ingesting pre-structured data formats (CSV, JSON arrays, database exports)
|
||||
- Inability to preserve the inherent structure when extracting tabular data from documents
|
||||
- Lack of efficient querying mechanisms for structured data patterns
|
||||
- Missing bridge between SQL-like queries and TrustGraph's graph queries
|
||||
|
||||
This specification addresses these gaps by introducing a structured data layer that complements TrustGraph's existing capabilities. By supporting structured data natively, TrustGraph can:
|
||||
- Serve as a unified platform for both structured and unstructured data analysis
|
||||
- Enable hybrid queries that span both graph relationships and tabular data
|
||||
- Provide familiar interfaces for users accustomed to working with structured data
|
||||
- Unlock new use cases in data integration and business intelligence
|
||||
|
||||
## Technical Design
|
||||
|
||||
### Architecture
|
||||
|
||||
The structured data integration requires the following technical components:
|
||||
|
||||
1. **NLP-to-Structured-Query Service**
|
||||
- Converts natural language questions into structured queries
|
||||
- Supports multiple query language targets (initially SQL-like syntax)
|
||||
- Integrates with existing TrustGraph NLP capabilities
|
||||
|
||||
Module: trustgraph-flow/trustgraph/query/nlp_query/cassandra
|
||||
|
||||
2. **Configuration Schema Support** ✅ **[COMPLETE]**
|
||||
- Extended configuration system to store structured data schemas
|
||||
- Support for defining table structures, field types, and relationships
|
||||
- Schema versioning and migration capabilities
|
||||
|
||||
3. **Object Extraction Module** ✅ **[COMPLETE]**
|
||||
- Enhanced knowledge extractor flow integration
|
||||
- Identifies and extracts structured objects from unstructured sources
|
||||
- Maintains provenance and confidence scores
|
||||
- Registers a config handler (example: trustgraph-flow/trustgraph/prompt/template/service.py) to receive config data and decode schema information
|
||||
- Receives objects and decodes them to ExtractedObject objects for delivery on the Pulsar queue
|
||||
- NOTE: There's existing code at `trustgraph-flow/trustgraph/extract/object/row/`. This was a previous attempt and will need to be majorly refactored as it doesn't conform to current APIs. Use it if it's useful, start from scratch if not.
|
||||
- Requires a command-line interface: `kg-extract-objects`
|
||||
|
||||
Module: trustgraph-flow/trustgraph/extract/kg/objects/
|
||||
|
||||
4. **Structured Store Writer Module** ✅ **[COMPLETE]**
|
||||
- Receives objects in ExtractedObject format from Pulsar queues
|
||||
- Initial implementation targeting Apache Cassandra as the structured data store
|
||||
- Handles dynamic table creation based on schemas encountered
|
||||
- Manages schema-to-Cassandra table mapping and data transformation
|
||||
- Provides batch and streaming write operations for performance optimization
|
||||
- No Pulsar outputs - this is a terminal service in the data flow
|
||||
|
||||
**Schema Handling**:
|
||||
- Monitors incoming ExtractedObject messages for schema references
|
||||
- When a new schema is encountered for the first time, automatically creates the corresponding Cassandra table
|
||||
- Maintains a cache of known schemas to avoid redundant table creation attempts
|
||||
- Should consider whether to receive schema definitions directly or rely on schema names in ExtractedObject messages
|
||||
|
||||
**Cassandra Table Mapping**:
|
||||
- Keyspace is named after the `user` field from ExtractedObject's Metadata
|
||||
- Table is named after the `schema_name` field from ExtractedObject
|
||||
- Collection from Metadata becomes part of the partition key to ensure:
|
||||
- Natural data distribution across Cassandra nodes
|
||||
- Efficient queries within a specific collection
|
||||
- Logical isolation between different data imports/sources
|
||||
- Primary key structure: `PRIMARY KEY ((collection, <schema_primary_key_fields>), <clustering_keys>)`
|
||||
- Collection is always the first component of the partition key
|
||||
- Schema-defined primary key fields follow as part of the composite partition key
|
||||
- This requires queries to specify the collection, ensuring predictable performance
|
||||
- Field definitions map to Cassandra columns with type conversions:
|
||||
- `string` → `text`
|
||||
- `integer` → `int` or `bigint` based on size hint
|
||||
- `float` → `float` or `double` based on precision needs
|
||||
- `boolean` → `boolean`
|
||||
- `timestamp` → `timestamp`
|
||||
- `enum` → `text` with application-level validation
|
||||
- Indexed fields create Cassandra secondary indexes (excluding fields already in the primary key)
|
||||
- Required fields are enforced at the application level (Cassandra doesn't support NOT NULL)
|
||||
|
||||
**Object Storage**:
|
||||
- Extracts values from ExtractedObject.values map
|
||||
- Performs type conversion and validation before insertion
|
||||
- Handles missing optional fields gracefully
|
||||
- Maintains metadata about object provenance (source document, confidence scores)
|
||||
- Supports idempotent writes to handle message replay scenarios
|
||||
|
||||
**Implementation Notes**:
|
||||
- Existing code at `trustgraph-flow/trustgraph/storage/objects/cassandra/` is outdated and doesn't comply with current APIs
|
||||
- Should reference `trustgraph-flow/trustgraph/storage/triples/cassandra` as an example of a working storage processor
|
||||
- Needs evaluation of existing code for any reusable components before deciding to refactor or rewrite
|
||||
|
||||
Module: trustgraph-flow/trustgraph/storage/objects/cassandra
|
||||
|
||||
5. **Structured Query Service**
|
||||
- Accepts structured queries in defined formats
|
||||
- Executes queries against the structured store
|
||||
- Returns objects matching query criteria
|
||||
- Supports pagination and result filtering
|
||||
|
||||
Module: trustgraph-flow/trustgraph/query/objects/cassandra
|
||||
|
||||
6. **Agent Tool Integration**
|
||||
- New tool class for agent frameworks
|
||||
- Enables agents to query structured data stores
|
||||
- Provides natural language and structured query interfaces
|
||||
- Integrates with existing agent decision-making processes
|
||||
|
||||
7. **Structured Data Ingestion Service**
|
||||
- Accepts structured data in multiple formats (JSON, CSV, XML)
|
||||
- Parses and validates incoming data against defined schemas
|
||||
- Converts data into normalized object streams
|
||||
- Emits objects to appropriate message queues for processing
|
||||
- Supports bulk uploads and streaming ingestion
|
||||
|
||||
Module: trustgraph-flow/trustgraph/decoding/structured
|
||||
|
||||
8. **Object Embedding Service**
|
||||
- Generates vector embeddings for structured objects
|
||||
- Enables semantic search across structured data
|
||||
- Supports hybrid search combining structured queries with semantic similarity
|
||||
- Integrates with existing vector stores
|
||||
|
||||
Module: trustgraph-flow/trustgraph/embeddings/object_embeddings/qdrant
|
||||
|
||||
### Data Models
|
||||
|
||||
#### Schema Storage Mechanism
|
||||
|
||||
Schemas are stored in TrustGraph's configuration system using the following structure:
|
||||
|
||||
- **Type**: `schema` (fixed value for all structured data schemas)
|
||||
- **Key**: The unique name/identifier of the schema (e.g., `customer_records`, `transaction_log`)
|
||||
- **Value**: JSON schema definition containing the structure
|
||||
|
||||
Example configuration entry:
|
||||
```
|
||||
Type: schema
|
||||
Key: customer_records
|
||||
Value: {
|
||||
"name": "customer_records",
|
||||
"description": "Customer information table",
|
||||
"fields": [
|
||||
{
|
||||
"name": "customer_id",
|
||||
"type": "string",
|
||||
"primary_key": true
|
||||
},
|
||||
{
|
||||
"name": "name",
|
||||
"type": "string",
|
||||
"required": true
|
||||
},
|
||||
{
|
||||
"name": "email",
|
||||
"type": "string",
|
||||
"required": true
|
||||
},
|
||||
{
|
||||
"name": "registration_date",
|
||||
"type": "timestamp"
|
||||
},
|
||||
{
|
||||
"name": "status",
|
||||
"type": "string",
|
||||
"enum": ["active", "inactive", "suspended"]
|
||||
}
|
||||
],
|
||||
"indexes": ["email", "registration_date"]
|
||||
}
|
||||
```
|
||||
|
||||
This approach allows:
|
||||
- Dynamic schema definition without code changes
|
||||
- Easy schema updates and versioning
|
||||
- Consistent integration with existing TrustGraph configuration management
|
||||
- Support for multiple schemas within a single deployment
|
||||
|
||||
### APIs
|
||||
|
||||
New APIs:
|
||||
- Pulsar schemas for above types
|
||||
- Pulsar interfaces in new flows
|
||||
- Need a means to specify schema types in flows so that flows know which
|
||||
schema types to load
|
||||
- APIs added to gateway and rev-gateway
|
||||
|
||||
Modified APIs:
|
||||
- Knowledge extraction endpoints - Add structured object output option
|
||||
- Agent endpoints - Add structured data tool support
|
||||
|
||||
### Implementation Details
|
||||
|
||||
Following existing conventions - these are just new processing modules.
|
||||
Everything is in the trustgraph-flow packages except for schema items
|
||||
in trustgraph-base.
|
||||
|
||||
Need some UI work in the Workbench to be able to demo / pilot this
|
||||
capability.
|
||||
|
||||
## Security Considerations
|
||||
|
||||
No extra considerations.
|
||||
|
||||
## Performance Considerations
|
||||
|
||||
Some questions around using Cassandra queries and indexes so that queries
|
||||
don't slow down.
|
||||
|
||||
## Testing Strategy
|
||||
|
||||
Use existing test strategy, will build unit, contract and integration tests.
|
||||
|
||||
## Migration Plan
|
||||
|
||||
None.
|
||||
|
||||
## Timeline
|
||||
|
||||
Not specified.
|
||||
|
||||
## Open Questions
|
||||
|
||||
- Can this be made to work with other store types? We're aiming to use
|
||||
interfaces which make modules which work with one store applicable to
|
||||
other stores.
|
||||
|
||||
## References
|
||||
|
||||
n/a.
|
||||
|
||||
|
|
@ -1,139 +0,0 @@
|
|||
# Structured Data Pulsar Schema Changes
|
||||
|
||||
## Overview
|
||||
|
||||
Based on the STRUCTURED_DATA.md specification, this document proposes the necessary Pulsar schema additions and modifications to support structured data capabilities in TrustGraph.
|
||||
|
||||
## Required Schema Changes
|
||||
|
||||
### 1. Core Schema Enhancements
|
||||
|
||||
#### Enhanced Field Definition
|
||||
The existing `Field` class in `core/primitives.py` needs additional properties:
|
||||
|
||||
```python
|
||||
class Field(Record):
|
||||
name = String()
|
||||
type = String() # int, string, long, bool, float, double, timestamp
|
||||
size = Integer()
|
||||
primary = Boolean()
|
||||
description = String()
|
||||
# NEW FIELDS:
|
||||
required = Boolean() # Whether field is required
|
||||
enum_values = Array(String()) # For enum type fields
|
||||
indexed = Boolean() # Whether field should be indexed
|
||||
```
|
||||
|
||||
### 2. New Knowledge Schemas
|
||||
|
||||
#### 2.1 Structured Data Submission
|
||||
New file: `knowledge/structured.py`
|
||||
|
||||
```python
|
||||
from pulsar.schema import Record, String, Bytes, Map
|
||||
from ..core.metadata import Metadata
|
||||
|
||||
class StructuredDataSubmission(Record):
|
||||
metadata = Metadata()
|
||||
format = String() # "json", "csv", "xml"
|
||||
schema_name = String() # Reference to schema in config
|
||||
data = Bytes() # Raw data to ingest
|
||||
options = Map(String()) # Format-specific options
|
||||
```
|
||||
|
||||
### 3. New Service Schemas
|
||||
|
||||
#### 3.1 NLP to Structured Query Service
|
||||
New file: `services/nlp_query.py`
|
||||
|
||||
```python
|
||||
from pulsar.schema import Record, String, Array, Map, Integer, Double
|
||||
from ..core.primitives import Error
|
||||
|
||||
class NLPToStructuredQueryRequest(Record):
|
||||
natural_language_query = String()
|
||||
max_results = Integer()
|
||||
context_hints = Map(String()) # Optional context for query generation
|
||||
|
||||
class NLPToStructuredQueryResponse(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()
|
||||
```
|
||||
|
||||
#### 3.2 Structured Query Service
|
||||
New file: `services/structured_query.py`
|
||||
|
||||
```python
|
||||
from pulsar.schema import Record, String, Map, Array
|
||||
from ..core.primitives import Error
|
||||
|
||||
class StructuredQueryRequest(Record):
|
||||
query = String() # GraphQL query
|
||||
variables = Map(String()) # GraphQL variables
|
||||
operation_name = String() # Optional operation name for multi-operation documents
|
||||
|
||||
class StructuredQueryResponse(Record):
|
||||
error = Error()
|
||||
data = String() # JSON-encoded GraphQL response data
|
||||
errors = Array(String()) # GraphQL errors if any
|
||||
```
|
||||
|
||||
#### 2.2 Object Extraction Output
|
||||
New file: `knowledge/object.py`
|
||||
|
||||
```python
|
||||
from pulsar.schema import Record, String, Map, Double
|
||||
from ..core.metadata import Metadata
|
||||
|
||||
class ExtractedObject(Record):
|
||||
metadata = Metadata()
|
||||
schema_name = String() # Which schema this object belongs to
|
||||
values = Map(String()) # Field name -> value
|
||||
confidence = Double()
|
||||
source_span = String() # Text span where object was found
|
||||
```
|
||||
|
||||
### 4. Enhanced Knowledge Schemas
|
||||
|
||||
#### 4.1 Object Embeddings Enhancement
|
||||
Update `knowledge/embeddings.py` to support structured object embeddings better:
|
||||
|
||||
```python
|
||||
class StructuredObjectEmbedding(Record):
|
||||
metadata = Metadata()
|
||||
vectors = Array(Array(Double()))
|
||||
schema_name = String()
|
||||
object_id = String() # Primary key value
|
||||
field_embeddings = Map(Array(Double())) # Per-field embeddings
|
||||
```
|
||||
|
||||
## Integration Points
|
||||
|
||||
### Flow Integration
|
||||
|
||||
The schemas will be used by new flow modules:
|
||||
- `trustgraph-flow/trustgraph/decoding/structured` - Uses StructuredDataSubmission
|
||||
- `trustgraph-flow/trustgraph/query/nlp_query/cassandra` - Uses NLP query schemas
|
||||
- `trustgraph-flow/trustgraph/query/objects/cassandra` - Uses structured query schemas
|
||||
- `trustgraph-flow/trustgraph/extract/object/row/` - Consumes Chunk, produces ExtractedObject
|
||||
- `trustgraph-flow/trustgraph/storage/objects/cassandra` - Uses Rows schema
|
||||
- `trustgraph-flow/trustgraph/embeddings/object_embeddings/qdrant` - Uses object embedding schemas
|
||||
|
||||
## Implementation Notes
|
||||
|
||||
1. **Schema Versioning**: Consider adding a `version` field to RowSchema for future migration support
|
||||
2. **Type System**: The `Field.type` should support all Cassandra native types
|
||||
3. **Batch Operations**: Most services should support both single and batch operations
|
||||
4. **Error Handling**: Consistent error reporting across all new services
|
||||
5. **Backwards Compatibility**: Existing schemas remain unchanged except for minor Field enhancements
|
||||
|
||||
## Next Steps
|
||||
|
||||
1. Implement schema files in the new structure
|
||||
2. Update existing services to recognize new schema types
|
||||
3. Implement flow modules that use these schemas
|
||||
4. Add gateway/rev-gateway endpoints for new services
|
||||
5. Create unit tests for schema validation
|
||||
570
docs/tech-specs/streaming-llm-responses.md
Normal file
570
docs/tech-specs/streaming-llm-responses.md
Normal file
|
|
@ -0,0 +1,570 @@
|
|||
# Streaming LLM Responses Technical Specification
|
||||
|
||||
## Overview
|
||||
|
||||
This specification describes the implementation of streaming support for LLM
|
||||
responses in TrustGraph. Streaming enables real-time delivery of generated
|
||||
tokens as they are produced by the LLM, rather than waiting for complete
|
||||
response generation.
|
||||
|
||||
This implementation supports the following use cases:
|
||||
|
||||
1. **Real-time User Interfaces**: Stream tokens to UI as they are generated,
|
||||
providing immediate visual feedback
|
||||
2. **Reduced Time-to-First-Token**: Users see output beginning immediately
|
||||
rather than waiting for full generation
|
||||
3. **Long Response Handling**: Handle very long outputs that might otherwise
|
||||
timeout or exceed memory limits
|
||||
4. **Interactive Applications**: Enable responsive chat and agent interfaces
|
||||
|
||||
## Goals
|
||||
|
||||
- **Backward Compatibility**: Existing non-streaming clients continue to work
|
||||
without modification
|
||||
- **Consistent API Design**: Streaming and non-streaming use the same schema
|
||||
patterns with minimal divergence
|
||||
- **Provider Flexibility**: Support streaming where available, graceful
|
||||
fallback where not
|
||||
- **Phased Rollout**: Incremental implementation to reduce risk
|
||||
- **End-to-End Support**: Streaming from LLM provider through to client
|
||||
applications via Pulsar, Gateway API, and Python API
|
||||
|
||||
## Background
|
||||
|
||||
### Current Architecture
|
||||
|
||||
The current LLM text completion flow operates as follows:
|
||||
|
||||
1. Client sends `TextCompletionRequest` with `system` and `prompt` fields
|
||||
2. LLM service processes the request and waits for complete generation
|
||||
3. Single `TextCompletionResponse` returned with complete `response` string
|
||||
|
||||
Current schema (`trustgraph-base/trustgraph/schema/services/llm.py`):
|
||||
|
||||
```python
|
||||
class TextCompletionRequest(Record):
|
||||
system = String()
|
||||
prompt = String()
|
||||
|
||||
class TextCompletionResponse(Record):
|
||||
error = Error()
|
||||
response = String()
|
||||
in_token = Integer()
|
||||
out_token = Integer()
|
||||
model = String()
|
||||
```
|
||||
|
||||
### Current Limitations
|
||||
|
||||
- **Latency**: Users must wait for complete generation before seeing any output
|
||||
- **Timeout Risk**: Long generations may exceed client timeout thresholds
|
||||
- **Poor UX**: No feedback during generation creates perception of slowness
|
||||
- **Resource Usage**: Full responses must be buffered in memory
|
||||
|
||||
This specification addresses these limitations by enabling incremental response
|
||||
delivery while maintaining full backward compatibility.
|
||||
|
||||
## Technical Design
|
||||
|
||||
### Phase 1: Infrastructure
|
||||
|
||||
Phase 1 establishes the foundation for streaming by modifying schemas, APIs,
|
||||
and CLI tools.
|
||||
|
||||
#### Schema Changes
|
||||
|
||||
##### LLM Schema (`trustgraph-base/trustgraph/schema/services/llm.py`)
|
||||
|
||||
**Request Changes:**
|
||||
|
||||
```python
|
||||
class TextCompletionRequest(Record):
|
||||
system = String()
|
||||
prompt = String()
|
||||
streaming = Boolean() # NEW: Default false for backward compatibility
|
||||
```
|
||||
|
||||
- `streaming`: When `true`, requests streaming response delivery
|
||||
- Default: `false` (existing behavior preserved)
|
||||
|
||||
**Response Changes:**
|
||||
|
||||
```python
|
||||
class TextCompletionResponse(Record):
|
||||
error = Error()
|
||||
response = String()
|
||||
in_token = Integer()
|
||||
out_token = Integer()
|
||||
model = String()
|
||||
end_of_stream = Boolean() # NEW: Indicates final message
|
||||
```
|
||||
|
||||
- `end_of_stream`: When `true`, indicates this is the final (or only) response
|
||||
- For non-streaming requests: Single response with `end_of_stream=true`
|
||||
- For streaming requests: Multiple responses, all with `end_of_stream=false`
|
||||
except the final one
|
||||
|
||||
##### Prompt Schema (`trustgraph-base/trustgraph/schema/services/prompt.py`)
|
||||
|
||||
The prompt service wraps text completion, so it mirrors the same pattern:
|
||||
|
||||
**Request Changes:**
|
||||
|
||||
```python
|
||||
class PromptRequest(Record):
|
||||
id = String()
|
||||
terms = Map(String())
|
||||
streaming = Boolean() # NEW: Default false
|
||||
```
|
||||
|
||||
**Response Changes:**
|
||||
|
||||
```python
|
||||
class PromptResponse(Record):
|
||||
error = Error()
|
||||
text = String()
|
||||
object = String()
|
||||
end_of_stream = Boolean() # NEW: Indicates final message
|
||||
```
|
||||
|
||||
#### Gateway API Changes
|
||||
|
||||
The Gateway API must expose streaming capabilities to HTTP/WebSocket clients.
|
||||
|
||||
**REST API Updates:**
|
||||
|
||||
- `POST /api/v1/text-completion`: Accept `streaming` parameter in request body
|
||||
- Response behavior depends on streaming flag:
|
||||
- `streaming=false`: Single JSON response (current behavior)
|
||||
- `streaming=true`: Server-Sent Events (SSE) stream or WebSocket messages
|
||||
|
||||
**Response Format (Streaming):**
|
||||
|
||||
Each streamed chunk follows the same schema structure:
|
||||
```json
|
||||
{
|
||||
"response": "partial text...",
|
||||
"end_of_stream": false,
|
||||
"model": "model-name"
|
||||
}
|
||||
```
|
||||
|
||||
Final chunk:
|
||||
```json
|
||||
{
|
||||
"response": "final text chunk",
|
||||
"end_of_stream": true,
|
||||
"in_token": 150,
|
||||
"out_token": 500,
|
||||
"model": "model-name"
|
||||
}
|
||||
```
|
||||
|
||||
#### Python API Changes
|
||||
|
||||
The Python client API must support both streaming and non-streaming modes
|
||||
while maintaining backward compatibility.
|
||||
|
||||
**LlmClient Updates** (`trustgraph-base/trustgraph/clients/llm_client.py`):
|
||||
|
||||
```python
|
||||
class LlmClient(BaseClient):
|
||||
def request(self, system, prompt, timeout=300, streaming=False):
|
||||
"""
|
||||
Non-streaming request (backward compatible).
|
||||
Returns complete response string.
|
||||
"""
|
||||
# Existing behavior when streaming=False
|
||||
|
||||
async def request_stream(self, system, prompt, timeout=300):
|
||||
"""
|
||||
Streaming request.
|
||||
Yields response chunks as they arrive.
|
||||
"""
|
||||
# New async generator method
|
||||
```
|
||||
|
||||
**PromptClient Updates** (`trustgraph-base/trustgraph/base/prompt_client.py`):
|
||||
|
||||
Similar pattern with `streaming` parameter and async generator variant.
|
||||
|
||||
#### CLI Tool Changes
|
||||
|
||||
**tg-invoke-llm** (`trustgraph-cli/trustgraph/cli/invoke_llm.py`):
|
||||
|
||||
```
|
||||
tg-invoke-llm [system] [prompt] [--no-streaming] [-u URL] [-f flow-id]
|
||||
```
|
||||
|
||||
- Streaming enabled by default for better interactive UX
|
||||
- `--no-streaming` flag disables streaming
|
||||
- When streaming: Output tokens to stdout as they arrive
|
||||
- When not streaming: Wait for complete response, then output
|
||||
|
||||
**tg-invoke-prompt** (`trustgraph-cli/trustgraph/cli/invoke_prompt.py`):
|
||||
|
||||
```
|
||||
tg-invoke-prompt [template-id] [var=value...] [--no-streaming] [-u URL] [-f flow-id]
|
||||
```
|
||||
|
||||
Same pattern as `tg-invoke-llm`.
|
||||
|
||||
#### LLM Service Base Class Changes
|
||||
|
||||
**LlmService** (`trustgraph-base/trustgraph/base/llm_service.py`):
|
||||
|
||||
```python
|
||||
class LlmService(FlowProcessor):
|
||||
async def on_request(self, msg, consumer, flow):
|
||||
request = msg.value()
|
||||
streaming = getattr(request, 'streaming', False)
|
||||
|
||||
if streaming and self.supports_streaming():
|
||||
async for chunk in self.generate_content_stream(...):
|
||||
await self.send_response(chunk, end_of_stream=False)
|
||||
await self.send_response(final_chunk, end_of_stream=True)
|
||||
else:
|
||||
response = await self.generate_content(...)
|
||||
await self.send_response(response, end_of_stream=True)
|
||||
|
||||
def supports_streaming(self):
|
||||
"""Override in subclass to indicate streaming support."""
|
||||
return False
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model, temperature):
|
||||
"""Override in subclass to implement streaming."""
|
||||
raise NotImplementedError()
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Phase 2: VertexAI Proof of Concept
|
||||
|
||||
Phase 2 implements streaming in a single provider (VertexAI) to validate the
|
||||
infrastructure and enable end-to-end testing.
|
||||
|
||||
#### VertexAI Implementation
|
||||
|
||||
**Module:** `trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py`
|
||||
|
||||
**Changes:**
|
||||
|
||||
1. Override `supports_streaming()` to return `True`
|
||||
2. Implement `generate_content_stream()` async generator
|
||||
3. Handle both Gemini and Claude models (via VertexAI Anthropic API)
|
||||
|
||||
**Gemini Streaming:**
|
||||
|
||||
```python
|
||||
async def generate_content_stream(self, system, prompt, model, temperature):
|
||||
model_instance = self.get_model(model, temperature)
|
||||
response = model_instance.generate_content(
|
||||
[system, prompt],
|
||||
stream=True # Enable streaming
|
||||
)
|
||||
for chunk in response:
|
||||
yield LlmChunk(
|
||||
text=chunk.text,
|
||||
in_token=None, # Available only in final chunk
|
||||
out_token=None,
|
||||
)
|
||||
# Final chunk includes token counts from response.usage_metadata
|
||||
```
|
||||
|
||||
**Claude (via VertexAI Anthropic) Streaming:**
|
||||
|
||||
```python
|
||||
async def generate_content_stream(self, system, prompt, model, temperature):
|
||||
with self.anthropic_client.messages.stream(...) as stream:
|
||||
for text in stream.text_stream:
|
||||
yield LlmChunk(text=text)
|
||||
# Token counts from stream.get_final_message()
|
||||
```
|
||||
|
||||
#### Testing
|
||||
|
||||
- Unit tests for streaming response assembly
|
||||
- Integration tests with VertexAI (Gemini and Claude)
|
||||
- End-to-end tests: CLI -> Gateway -> Pulsar -> VertexAI -> back
|
||||
- Backward compatibility tests: Non-streaming requests still work
|
||||
|
||||
---
|
||||
|
||||
### Phase 3: All LLM Providers
|
||||
|
||||
Phase 3 extends streaming support to all LLM providers in the system.
|
||||
|
||||
#### Provider Implementation Status
|
||||
|
||||
Each provider must either:
|
||||
1. **Full Streaming Support**: Implement `generate_content_stream()`
|
||||
2. **Compatibility Mode**: Handle the `end_of_stream` flag correctly
|
||||
(return single response with `end_of_stream=true`)
|
||||
|
||||
| Provider | Package | Streaming Support |
|
||||
|----------|---------|-------------------|
|
||||
| OpenAI | trustgraph-flow | Full (native streaming API) |
|
||||
| Claude/Anthropic | trustgraph-flow | Full (native streaming API) |
|
||||
| Ollama | trustgraph-flow | Full (native streaming API) |
|
||||
| Cohere | trustgraph-flow | Full (native streaming API) |
|
||||
| Mistral | trustgraph-flow | Full (native streaming API) |
|
||||
| Azure OpenAI | trustgraph-flow | Full (native streaming API) |
|
||||
| Google AI Studio | trustgraph-flow | Full (native streaming API) |
|
||||
| VertexAI | trustgraph-vertexai | Full (Phase 2) |
|
||||
| Bedrock | trustgraph-bedrock | Full (native streaming API) |
|
||||
| LM Studio | trustgraph-flow | Full (OpenAI-compatible) |
|
||||
| LlamaFile | trustgraph-flow | Full (OpenAI-compatible) |
|
||||
| vLLM | trustgraph-flow | Full (OpenAI-compatible) |
|
||||
| TGI | trustgraph-flow | TBD |
|
||||
| Azure | trustgraph-flow | TBD |
|
||||
|
||||
#### Implementation Pattern
|
||||
|
||||
For OpenAI-compatible providers (OpenAI, LM Studio, LlamaFile, vLLM):
|
||||
|
||||
```python
|
||||
async def generate_content_stream(self, system, prompt, model, temperature):
|
||||
response = await self.client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[
|
||||
{"role": "system", "content": system},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=temperature,
|
||||
stream=True
|
||||
)
|
||||
async for chunk in response:
|
||||
if chunk.choices[0].delta.content:
|
||||
yield LlmChunk(text=chunk.choices[0].delta.content)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Phase 4: Agent API
|
||||
|
||||
Phase 4 extends streaming to the Agent API. This is more complex because the
|
||||
Agent API is already multi-message by nature (thought → action → observation
|
||||
→ repeat → final answer).
|
||||
|
||||
#### Current Agent Schema
|
||||
|
||||
```python
|
||||
class AgentStep(Record):
|
||||
thought = String()
|
||||
action = String()
|
||||
arguments = Map(String())
|
||||
observation = String()
|
||||
user = String()
|
||||
|
||||
class AgentRequest(Record):
|
||||
question = String()
|
||||
state = String()
|
||||
group = Array(String())
|
||||
history = Array(AgentStep())
|
||||
user = String()
|
||||
|
||||
class AgentResponse(Record):
|
||||
answer = String()
|
||||
error = Error()
|
||||
thought = String()
|
||||
observation = String()
|
||||
```
|
||||
|
||||
#### Proposed Agent Schema Changes
|
||||
|
||||
**Request Changes:**
|
||||
|
||||
```python
|
||||
class AgentRequest(Record):
|
||||
question = String()
|
||||
state = String()
|
||||
group = Array(String())
|
||||
history = Array(AgentStep())
|
||||
user = String()
|
||||
streaming = Boolean() # NEW: Default false
|
||||
```
|
||||
|
||||
**Response Changes:**
|
||||
|
||||
The agent produces multiple types of output during its reasoning cycle:
|
||||
- Thoughts (reasoning)
|
||||
- Actions (tool calls)
|
||||
- Observations (tool results)
|
||||
- Answer (final response)
|
||||
- Errors
|
||||
|
||||
Since `chunk_type` identifies what kind of content is being sent, the separate
|
||||
`answer`, `error`, `thought`, and `observation` fields can be collapsed into
|
||||
a single `content` field:
|
||||
|
||||
```python
|
||||
class AgentResponse(Record):
|
||||
chunk_type = String() # "thought", "action", "observation", "answer", "error"
|
||||
content = String() # The actual content (interpretation depends on chunk_type)
|
||||
end_of_message = Boolean() # Current thought/action/observation/answer is complete
|
||||
end_of_dialog = Boolean() # Entire agent dialog is complete
|
||||
```
|
||||
|
||||
**Field Semantics:**
|
||||
|
||||
- `chunk_type`: Indicates what type of content is in the `content` field
|
||||
- `"thought"`: Agent reasoning/thinking
|
||||
- `"action"`: Tool/action being invoked
|
||||
- `"observation"`: Result from tool execution
|
||||
- `"answer"`: Final answer to the user's question
|
||||
- `"error"`: Error message
|
||||
|
||||
- `content`: The actual streamed content, interpreted based on `chunk_type`
|
||||
|
||||
- `end_of_message`: When `true`, the current chunk type is complete
|
||||
- Example: All tokens for the current thought have been sent
|
||||
- Allows clients to know when to move to the next stage
|
||||
|
||||
- `end_of_dialog`: When `true`, the entire agent interaction is complete
|
||||
- This is the final message in the stream
|
||||
|
||||
#### Agent Streaming Behavior
|
||||
|
||||
When `streaming=true`:
|
||||
|
||||
1. **Thought streaming**:
|
||||
- Multiple chunks with `chunk_type="thought"`, `end_of_message=false`
|
||||
- Final thought chunk has `end_of_message=true`
|
||||
2. **Action notification**:
|
||||
- Single chunk with `chunk_type="action"`, `end_of_message=true`
|
||||
3. **Observation**:
|
||||
- Chunk(s) with `chunk_type="observation"`, final has `end_of_message=true`
|
||||
4. **Repeat** steps 1-3 as the agent reasons
|
||||
5. **Final answer**:
|
||||
- `chunk_type="answer"` with the final response in `content`
|
||||
- Last chunk has `end_of_message=true`, `end_of_dialog=true`
|
||||
|
||||
**Example Stream Sequence:**
|
||||
|
||||
```
|
||||
{chunk_type: "thought", content: "I need to", end_of_message: false, end_of_dialog: false}
|
||||
{chunk_type: "thought", content: " search for...", end_of_message: true, end_of_dialog: false}
|
||||
{chunk_type: "action", content: "search", end_of_message: true, end_of_dialog: false}
|
||||
{chunk_type: "observation", content: "Found: ...", end_of_message: true, end_of_dialog: false}
|
||||
{chunk_type: "thought", content: "Based on this", end_of_message: false, end_of_dialog: false}
|
||||
{chunk_type: "thought", content: " I can answer...", end_of_message: true, end_of_dialog: false}
|
||||
{chunk_type: "answer", content: "The answer is...", end_of_message: true, end_of_dialog: true}
|
||||
```
|
||||
|
||||
When `streaming=false`:
|
||||
- Current behavior preserved
|
||||
- Single response with complete answer
|
||||
- `end_of_message=true`, `end_of_dialog=true`
|
||||
|
||||
#### Gateway and Python API
|
||||
|
||||
- Gateway: New SSE/WebSocket endpoint for agent streaming
|
||||
- Python API: New `agent_stream()` async generator method
|
||||
|
||||
---
|
||||
|
||||
## Security Considerations
|
||||
|
||||
- **No new attack surface**: Streaming uses same authentication/authorization
|
||||
- **Rate limiting**: Apply per-token or per-chunk rate limits if needed
|
||||
- **Connection handling**: Properly terminate streams on client disconnect
|
||||
- **Timeout management**: Streaming requests need appropriate timeout handling
|
||||
|
||||
## Performance Considerations
|
||||
|
||||
- **Memory**: Streaming reduces peak memory usage (no full response buffering)
|
||||
- **Latency**: Time-to-first-token significantly reduced
|
||||
- **Connection overhead**: SSE/WebSocket connections have keep-alive overhead
|
||||
- **Pulsar throughput**: Multiple small messages vs. single large message
|
||||
tradeoff
|
||||
|
||||
## Testing Strategy
|
||||
|
||||
### Unit Tests
|
||||
- Schema serialization/deserialization with new fields
|
||||
- Backward compatibility (missing fields use defaults)
|
||||
- Chunk assembly logic
|
||||
|
||||
### Integration Tests
|
||||
- Each LLM provider's streaming implementation
|
||||
- Gateway API streaming endpoints
|
||||
- Python client streaming methods
|
||||
|
||||
### End-to-End Tests
|
||||
- CLI tool streaming output
|
||||
- Full flow: Client → Gateway → Pulsar → LLM → back
|
||||
- Mixed streaming/non-streaming workloads
|
||||
|
||||
### Backward Compatibility Tests
|
||||
- Existing clients work without modification
|
||||
- Non-streaming requests behave identically
|
||||
|
||||
## Migration Plan
|
||||
|
||||
### Phase 1: Infrastructure
|
||||
- Deploy schema changes (backward compatible)
|
||||
- Deploy Gateway API updates
|
||||
- Deploy Python API updates
|
||||
- Release CLI tool updates
|
||||
|
||||
### Phase 2: VertexAI
|
||||
- Deploy VertexAI streaming implementation
|
||||
- Validate with test workloads
|
||||
|
||||
### Phase 3: All Providers
|
||||
- Roll out provider updates incrementally
|
||||
- Monitor for issues
|
||||
|
||||
### Phase 4: Agent API
|
||||
- Deploy agent schema changes
|
||||
- Deploy agent streaming implementation
|
||||
- Update documentation
|
||||
|
||||
## Timeline
|
||||
|
||||
| Phase | Description | Dependencies |
|
||||
|-------|-------------|--------------|
|
||||
| Phase 1 | Infrastructure | None |
|
||||
| Phase 2 | VertexAI PoC | Phase 1 |
|
||||
| Phase 3 | All Providers | Phase 2 |
|
||||
| Phase 4 | Agent API | Phase 3 |
|
||||
|
||||
## Design Decisions
|
||||
|
||||
The following questions were resolved during specification:
|
||||
|
||||
1. **Token Counts in Streaming**: Token counts are deltas, not running totals.
|
||||
Consumers can sum them if needed. This matches how most providers report
|
||||
usage and simplifies the implementation.
|
||||
|
||||
2. **Error Handling in Streams**: If an error occurs, the `error` field is
|
||||
populated and no other fields are needed. An error is always the final
|
||||
communication - no subsequent messages are permitted or expected after
|
||||
an error. For LLM/Prompt streams, `end_of_stream=true`. For Agent streams,
|
||||
`chunk_type="error"` with `end_of_dialog=true`.
|
||||
|
||||
3. **Partial Response Recovery**: The messaging protocol (Pulsar) is resilient,
|
||||
so message-level retry is not needed. If a client loses track of the stream
|
||||
or disconnects, it must retry the full request from scratch.
|
||||
|
||||
4. **Prompt Service Streaming**: Streaming is only supported for text (`text`)
|
||||
responses, not structured (`object`) responses. The prompt service knows at
|
||||
the outset whether the output will be JSON or text based on the prompt
|
||||
template. If a streaming request is made for a JSON-output prompt, the
|
||||
service should either:
|
||||
- Return the complete JSON in a single response with `end_of_stream=true`, or
|
||||
- Reject the streaming request with an error
|
||||
|
||||
## Open Questions
|
||||
|
||||
None at this time.
|
||||
|
||||
## References
|
||||
|
||||
- Current LLM schema: `trustgraph-base/trustgraph/schema/services/llm.py`
|
||||
- Current prompt schema: `trustgraph-base/trustgraph/schema/services/prompt.py`
|
||||
- Current agent schema: `trustgraph-base/trustgraph/schema/services/agent.py`
|
||||
- LLM service base: `trustgraph-base/trustgraph/base/llm_service.py`
|
||||
- VertexAI provider: `trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py`
|
||||
- Gateway API: `trustgraph-base/trustgraph/api/`
|
||||
- CLI tools: `trustgraph-cli/trustgraph/cli/`
|
||||
|
|
@ -282,10 +282,11 @@ class TestTextCompletionIntegration:
|
|||
# Assert
|
||||
# Verify OpenAI API call parameters
|
||||
call_args = mock_openai_client.chat.completions.create.call_args
|
||||
assert call_args.kwargs['response_format'] == {"type": "text"}
|
||||
assert call_args.kwargs['top_p'] == 1
|
||||
assert call_args.kwargs['frequency_penalty'] == 0
|
||||
assert call_args.kwargs['presence_penalty'] == 0
|
||||
# Note: response_format, top_p, frequency_penalty, and presence_penalty
|
||||
# were removed in #561 as unnecessary parameters
|
||||
assert 'model' in call_args.kwargs
|
||||
assert 'temperature' in call_args.kwargs
|
||||
assert 'max_tokens' in call_args.kwargs
|
||||
|
||||
# Verify result structure
|
||||
assert hasattr(result, 'text')
|
||||
|
|
@ -362,9 +363,8 @@ class TestTextCompletionIntegration:
|
|||
assert call_args.kwargs['model'] == "gpt-4"
|
||||
assert call_args.kwargs['temperature'] == 0.8
|
||||
assert call_args.kwargs['max_tokens'] == 2048
|
||||
assert call_args.kwargs['top_p'] == 1
|
||||
assert call_args.kwargs['frequency_penalty'] == 0
|
||||
assert call_args.kwargs['presence_penalty'] == 0
|
||||
# Note: top_p, frequency_penalty, and presence_penalty
|
||||
# were removed in #561 as unnecessary parameters
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.slow
|
||||
|
|
|
|||
|
|
@ -5,6 +5,9 @@ Tests for Pinecone document embeddings query service
|
|||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
# Skip all tests in this module due to missing Pinecone dependency
|
||||
pytest.skip("Pinecone library missing protoc_gen_openapiv2 dependency", allow_module_level=True)
|
||||
|
||||
from trustgraph.query.doc_embeddings.pinecone.service import Processor
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -5,6 +5,9 @@ Tests for Pinecone graph embeddings query service
|
|||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
# Skip all tests in this module due to missing Pinecone dependency
|
||||
pytest.skip("Pinecone library missing protoc_gen_openapiv2 dependency", allow_module_level=True)
|
||||
|
||||
from trustgraph.query.graph_embeddings.pinecone.service import Processor
|
||||
from trustgraph.schema import Value
|
||||
|
||||
|
|
|
|||
|
|
@ -6,6 +6,9 @@ import pytest
|
|||
from unittest.mock import MagicMock, patch
|
||||
import uuid
|
||||
|
||||
# Skip all tests in this module due to missing Pinecone dependency
|
||||
pytest.skip("Pinecone library missing protoc_gen_openapiv2 dependency", allow_module_level=True)
|
||||
|
||||
from trustgraph.storage.doc_embeddings.pinecone.write import Processor
|
||||
from trustgraph.schema import ChunkEmbeddings
|
||||
|
||||
|
|
|
|||
|
|
@ -6,6 +6,9 @@ import pytest
|
|||
from unittest.mock import MagicMock, patch
|
||||
import uuid
|
||||
|
||||
# Skip all tests in this module due to missing Pinecone dependency
|
||||
pytest.skip("Pinecone library missing protoc_gen_openapiv2 dependency", allow_module_level=True)
|
||||
|
||||
from trustgraph.storage.graph_embeddings.pinecone.write import Processor
|
||||
from trustgraph.schema import EntityEmbeddings, Value
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
model = flow("model")
|
||||
temperature = flow("temperature")
|
||||
# Check if streaming is requested and supported
|
||||
streaming = getattr(request, 'streaming', False)
|
||||
|
||||
response = await self.generate_content(
|
||||
request.system, request.prompt, model, temperature
|
||||
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):
|
||||
|
||||
|
|
|
|||
|
|
@ -1,30 +1,75 @@
|
|||
|
||||
import json
|
||||
import asyncio
|
||||
|
||||
from . request_response_spec import RequestResponse, RequestResponseSpec
|
||||
from .. schema import PromptRequest, PromptResponse
|
||||
|
||||
class PromptClient(RequestResponse):
|
||||
|
||||
async def prompt(self, id, variables, timeout=600):
|
||||
async def prompt(self, id, variables, timeout=600, streaming=False, chunk_callback=None):
|
||||
|
||||
resp = await self.request(
|
||||
PromptRequest(
|
||||
id = id,
|
||||
terms = {
|
||||
k: json.dumps(v)
|
||||
for k, v in variables.items()
|
||||
}
|
||||
),
|
||||
timeout=timeout
|
||||
)
|
||||
if not streaming:
|
||||
# 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.error:
|
||||
raise RuntimeError(resp.error.message)
|
||||
|
||||
if resp.text: return resp.text
|
||||
if resp.text: return resp.text
|
||||
|
||||
return json.loads(resp.object)
|
||||
return json.loads(resp.object)
|
||||
|
||||
else:
|
||||
# Streaming path - collect all chunks
|
||||
full_text = ""
|
||||
full_object = None
|
||||
|
||||
async def collect_chunks(resp):
|
||||
nonlocal full_text, full_object
|
||||
|
||||
if resp.error:
|
||||
raise RuntimeError(resp.error.message)
|
||||
|
||||
if resp.text:
|
||||
full_text += resp.text
|
||||
# Call chunk callback if provided
|
||||
if chunk_callback:
|
||||
if asyncio.iscoroutinefunction(chunk_callback):
|
||||
await chunk_callback(resp.text)
|
||||
else:
|
||||
chunk_callback(resp.text)
|
||||
elif resp.object:
|
||||
full_object = resp.object
|
||||
|
||||
return getattr(resp, 'end_of_stream', False)
|
||||
|
||||
await self.request(
|
||||
PromptRequest(
|
||||
id = id,
|
||||
terms = {
|
||||
k: json.dumps(v)
|
||||
for k, v in variables.items()
|
||||
},
|
||||
streaming = True
|
||||
),
|
||||
recipient=collect_chunks,
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
if full_text: return full_text
|
||||
|
||||
return json.loads(full_object)
|
||||
|
||||
async def extract_definitions(self, text, timeout=600):
|
||||
return await self.prompt(
|
||||
|
|
@ -70,11 +115,13 @@ class PromptClient(RequestResponse):
|
|||
timeout = timeout,
|
||||
)
|
||||
|
||||
async def agent_react(self, variables, timeout=600):
|
||||
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):
|
||||
|
|
|
|||
|
|
@ -3,18 +3,45 @@ from . request_response_spec import RequestResponse, RequestResponseSpec
|
|||
from .. schema import TextCompletionRequest, TextCompletionResponse
|
||||
|
||||
class TextCompletionClient(RequestResponse):
|
||||
async def text_completion(self, system, prompt, timeout=600):
|
||||
resp = await self.request(
|
||||
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
|
||||
system = system, prompt = prompt, streaming = True
|
||||
),
|
||||
recipient=collect_chunks,
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
if resp.error:
|
||||
raise RuntimeError(resp.error.message)
|
||||
|
||||
return resp.response
|
||||
return full_response
|
||||
|
||||
class TextCompletionClientSpec(RequestResponseSpec):
|
||||
def __init__(
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ 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
|
||||
|
|
@ -37,8 +38,68 @@ class LlmClient(BaseClient):
|
|||
output_schema=TextCompletionResponse,
|
||||
)
|
||||
|
||||
def request(self, system, prompt, timeout=300):
|
||||
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, timeout=timeout
|
||||
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")
|
||||
|
||||
|
|
|
|||
|
|
@ -12,16 +12,18 @@ class AgentRequestTranslator(MessageTranslator):
|
|||
state=data.get("state", None),
|
||||
group=data.get("group", None),
|
||||
history=data.get("history", []),
|
||||
user=data.get("user", "trustgraph")
|
||||
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
|
||||
"user": obj.user,
|
||||
"streaming": getattr(obj, "streaming", False)
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -33,14 +35,36 @@ class AgentResponseTranslator(MessageTranslator):
|
|||
|
||||
def from_pulsar(self, obj: AgentResponse) -> Dict[str, Any]:
|
||||
result = {}
|
||||
if obj.answer:
|
||||
result["answer"] = obj.answer
|
||||
if obj.thought:
|
||||
result["thought"] = obj.thought
|
||||
if obj.observation:
|
||||
result["observation"] = obj.observation
|
||||
|
||||
# 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)"""
|
||||
return self.from_pulsar(obj), (obj.answer is not None)
|
||||
# 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
|
||||
|
|
@ -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
|
||||
|
|
@ -5,11 +5,12 @@ from .base import MessageTranslator
|
|||
|
||||
class TextCompletionRequestTranslator(MessageTranslator):
|
||||
"""Translator for TextCompletionRequest schema objects"""
|
||||
|
||||
|
||||
def to_pulsar(self, data: Dict[str, Any]) -> TextCompletionRequest:
|
||||
return TextCompletionRequest(
|
||||
system=data["system"],
|
||||
prompt=data["prompt"]
|
||||
prompt=data["prompt"],
|
||||
streaming=data.get("streaming", False)
|
||||
)
|
||||
|
||||
def from_pulsar(self, obj: TextCompletionRequest) -> Dict[str, Any]:
|
||||
|
|
@ -39,4 +40,6 @@ class TextCompletionResponseTranslator(MessageTranslator):
|
|||
|
||||
def from_response_with_completion(self, obj: TextCompletionResponse) -> 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
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
|
||||
from pulsar.schema import Record, String, Array, Map
|
||||
from pulsar.schema import Record, String, Array, Map, Boolean
|
||||
|
||||
from ..core.topic import topic
|
||||
from ..core.primitives import Error
|
||||
|
|
@ -21,8 +21,16 @@ class AgentRequest(Record):
|
|||
group = Array(String())
|
||||
history = Array(AgentStep())
|
||||
user = String() # User context for multi-tenancy
|
||||
streaming = Boolean() # NEW: Enable streaming response delivery (default false)
|
||||
|
||||
class AgentResponse(Record):
|
||||
# Streaming-first design
|
||||
chunk_type = String() # "thought", "action", "observation", "answer", "error"
|
||||
content = String() # The actual content (interpretation depends on chunk_type)
|
||||
end_of_message = Boolean() # Current chunk type (thought/action/etc.) is complete
|
||||
end_of_dialog = Boolean() # Entire agent dialog is complete
|
||||
|
||||
# Legacy fields (deprecated but kept for backward compatibility)
|
||||
answer = String()
|
||||
error = Error()
|
||||
thought = String()
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
|
||||
from pulsar.schema import Record, String, Array, Double, Integer
|
||||
from pulsar.schema import Record, String, Array, Double, Integer, Boolean
|
||||
|
||||
from ..core.topic import topic
|
||||
from ..core.primitives import Error
|
||||
|
|
@ -11,6 +11,7 @@ from ..core.primitives import Error
|
|||
class TextCompletionRequest(Record):
|
||||
system = String()
|
||||
prompt = String()
|
||||
streaming = Boolean() # Default false for backward compatibility
|
||||
|
||||
class TextCompletionResponse(Record):
|
||||
error = Error()
|
||||
|
|
@ -18,6 +19,7 @@ class TextCompletionResponse(Record):
|
|||
in_token = Integer()
|
||||
out_token = Integer()
|
||||
model = String()
|
||||
end_of_stream = Boolean() # Indicates final message in stream
|
||||
|
||||
############################################################################
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from pulsar.schema import Record, String, Map
|
||||
from pulsar.schema import Record, String, Map, Boolean
|
||||
|
||||
from ..core.primitives import Error
|
||||
from ..core.topic import topic
|
||||
|
|
@ -24,6 +24,9 @@ class PromptRequest(Record):
|
|||
# JSON encoded values
|
||||
terms = Map(String())
|
||||
|
||||
# Streaming support (default false for backward compatibility)
|
||||
streaming = Boolean()
|
||||
|
||||
class PromptResponse(Record):
|
||||
|
||||
# Error case
|
||||
|
|
@ -35,4 +38,7 @@ class PromptResponse(Record):
|
|||
# JSON encoded
|
||||
object = String()
|
||||
|
||||
# Indicates final message in stream
|
||||
end_of_stream = Boolean()
|
||||
|
||||
############################################################################
|
||||
|
|
@ -11,7 +11,7 @@ import enum
|
|||
import logging
|
||||
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import LlmService, LlmResult
|
||||
from .... base import LlmService, LlmResult, LlmChunk
|
||||
|
||||
# Module logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -52,6 +52,8 @@ class ModelHandler:
|
|||
raise RuntimeError("format_request not implemented")
|
||||
def decode_response(self, response):
|
||||
raise RuntimeError("format_request not implemented")
|
||||
def decode_stream_chunk(self, chunk):
|
||||
raise RuntimeError("decode_stream_chunk not implemented")
|
||||
|
||||
class Mistral(ModelHandler):
|
||||
def __init__(self):
|
||||
|
|
@ -68,6 +70,11 @@ class Mistral(ModelHandler):
|
|||
def decode_response(self, response):
|
||||
response_body = json.loads(response.get("body").read())
|
||||
return response_body['outputs'][0]['text']
|
||||
def decode_stream_chunk(self, chunk):
|
||||
chunk_obj = json.loads(chunk.get('chunk').get('bytes').decode())
|
||||
if 'outputs' in chunk_obj and len(chunk_obj['outputs']) > 0:
|
||||
return chunk_obj['outputs'][0].get('text', '')
|
||||
return ''
|
||||
|
||||
# Llama 3
|
||||
class Meta(ModelHandler):
|
||||
|
|
@ -83,6 +90,9 @@ class Meta(ModelHandler):
|
|||
def decode_response(self, response):
|
||||
model_response = json.loads(response["body"].read())
|
||||
return model_response["generation"]
|
||||
def decode_stream_chunk(self, chunk):
|
||||
chunk_obj = json.loads(chunk.get('chunk').get('bytes').decode())
|
||||
return chunk_obj.get('generation', '')
|
||||
|
||||
class Anthropic(ModelHandler):
|
||||
def __init__(self):
|
||||
|
|
@ -108,6 +118,12 @@ class Anthropic(ModelHandler):
|
|||
def decode_response(self, response):
|
||||
model_response = json.loads(response["body"].read())
|
||||
return model_response['content'][0]['text']
|
||||
def decode_stream_chunk(self, chunk):
|
||||
chunk_obj = json.loads(chunk.get('chunk').get('bytes').decode())
|
||||
if chunk_obj.get('type') == 'content_block_delta':
|
||||
if 'delta' in chunk_obj and 'text' in chunk_obj['delta']:
|
||||
return chunk_obj['delta']['text']
|
||||
return ''
|
||||
|
||||
class Ai21(ModelHandler):
|
||||
def __init__(self):
|
||||
|
|
@ -129,6 +145,12 @@ class Ai21(ModelHandler):
|
|||
content_str = content.decode('utf-8')
|
||||
content_json = json.loads(content_str)
|
||||
return content_json['choices'][0]['message']['content']
|
||||
def decode_stream_chunk(self, chunk):
|
||||
chunk_obj = json.loads(chunk.get('chunk').get('bytes').decode())
|
||||
if 'choices' in chunk_obj and len(chunk_obj['choices']) > 0:
|
||||
delta = chunk_obj['choices'][0].get('delta', {})
|
||||
return delta.get('content', '')
|
||||
return ''
|
||||
|
||||
class Cohere(ModelHandler):
|
||||
def encode_request(self, system, prompt):
|
||||
|
|
@ -142,6 +164,9 @@ class Cohere(ModelHandler):
|
|||
content_str = content.decode('utf-8')
|
||||
content_json = json.loads(content_str)
|
||||
return content_json['text']
|
||||
def decode_stream_chunk(self, chunk):
|
||||
chunk_obj = json.loads(chunk.get('chunk').get('bytes').decode())
|
||||
return chunk_obj.get('text', '')
|
||||
|
||||
Default=Mistral
|
||||
|
||||
|
|
@ -309,6 +334,78 @@ class Processor(LlmService):
|
|||
logger.error(f"Bedrock LLM exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
def supports_streaming(self):
|
||||
"""Bedrock supports streaming"""
|
||||
return True
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
||||
"""Stream content generation from Bedrock"""
|
||||
model_name = model or self.default_model
|
||||
effective_temperature = temperature if temperature is not None else self.temperature
|
||||
|
||||
logger.debug(f"Using model (streaming): {model_name}")
|
||||
logger.debug(f"Using temperature: {effective_temperature}")
|
||||
|
||||
try:
|
||||
variant = self._get_or_create_variant(model_name, effective_temperature)
|
||||
promptbody = variant.encode_request(system, prompt)
|
||||
|
||||
accept = 'application/json'
|
||||
contentType = 'application/json'
|
||||
|
||||
response = self.bedrock.invoke_model_with_response_stream(
|
||||
body=promptbody,
|
||||
modelId=model_name,
|
||||
accept=accept,
|
||||
contentType=contentType
|
||||
)
|
||||
|
||||
total_input_tokens = 0
|
||||
total_output_tokens = 0
|
||||
|
||||
stream = response.get('body')
|
||||
if stream:
|
||||
for event in stream:
|
||||
chunk = event.get('chunk')
|
||||
if chunk:
|
||||
# Decode the chunk text
|
||||
text = variant.decode_stream_chunk(event)
|
||||
if text:
|
||||
yield LlmChunk(
|
||||
text=text,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=False
|
||||
)
|
||||
|
||||
# Try to extract metadata from the event
|
||||
metadata = event.get('metadata')
|
||||
if metadata:
|
||||
usage = metadata.get('usage')
|
||||
if usage:
|
||||
total_input_tokens = usage.get('inputTokens', 0)
|
||||
total_output_tokens = usage.get('outputTokens', 0)
|
||||
|
||||
# Send final chunk with token counts
|
||||
yield LlmChunk(
|
||||
text="",
|
||||
in_token=total_input_tokens,
|
||||
out_token=total_output_tokens,
|
||||
model=model_name,
|
||||
is_final=True
|
||||
)
|
||||
|
||||
logger.debug("Streaming complete")
|
||||
|
||||
except self.bedrock.exceptions.ThrottlingException as e:
|
||||
logger.warning(f"Hit rate limit during streaming: {e}")
|
||||
raise TooManyRequests()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Bedrock streaming exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ def output(text, prefix="> ", width=78):
|
|||
|
||||
async def question(
|
||||
url, question, flow_id, user, collection,
|
||||
plan=None, state=None, group=None, verbose=False
|
||||
plan=None, state=None, group=None, verbose=False, streaming=True
|
||||
):
|
||||
|
||||
if not url.endswith("/"):
|
||||
|
|
@ -62,16 +62,17 @@ async def question(
|
|||
"request": {
|
||||
"question": question,
|
||||
"user": user,
|
||||
"history": []
|
||||
"history": [],
|
||||
"streaming": streaming
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
# Only add optional fields if they have values
|
||||
if state is not None:
|
||||
req["request"]["state"] = state
|
||||
if group is not None:
|
||||
req["request"]["group"] = group
|
||||
|
||||
|
||||
req = json.dumps(req)
|
||||
|
||||
await ws.send(req)
|
||||
|
|
@ -89,14 +90,34 @@ async def question(
|
|||
print("Ignore message")
|
||||
continue
|
||||
|
||||
if "thought" in obj["response"]:
|
||||
think(obj["response"]["thought"])
|
||||
response = obj["response"]
|
||||
|
||||
if "observation" in obj["response"]:
|
||||
observe(obj["response"]["observation"])
|
||||
# Handle streaming format (new format with chunk_type)
|
||||
if "chunk_type" in response:
|
||||
chunk_type = response["chunk_type"]
|
||||
content = response.get("content", "")
|
||||
|
||||
if "answer" in obj["response"]:
|
||||
print(obj["response"]["answer"])
|
||||
if chunk_type == "thought":
|
||||
think(content)
|
||||
elif chunk_type == "observation":
|
||||
observe(content)
|
||||
elif chunk_type == "answer":
|
||||
print(content)
|
||||
elif chunk_type == "error":
|
||||
raise RuntimeError(content)
|
||||
else:
|
||||
# Handle legacy format (backward compatibility)
|
||||
if "thought" in response:
|
||||
think(response["thought"])
|
||||
|
||||
if "observation" in response:
|
||||
observe(response["observation"])
|
||||
|
||||
if "answer" in response:
|
||||
print(response["answer"])
|
||||
|
||||
if "error" in response:
|
||||
raise RuntimeError(response["error"])
|
||||
|
||||
if obj["complete"]: break
|
||||
|
||||
|
|
@ -161,6 +182,12 @@ def main():
|
|||
help=f'Output thinking/observations'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--no-streaming',
|
||||
action="store_true",
|
||||
help=f'Disable streaming (use legacy mode)'
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
|
|
@ -176,6 +203,7 @@ def main():
|
|||
state = args.state,
|
||||
group = args.group,
|
||||
verbose = args.verbose,
|
||||
streaming = not args.no_streaming,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -6,17 +6,63 @@ and user prompt. Both arguments are required.
|
|||
import argparse
|
||||
import os
|
||||
import json
|
||||
from trustgraph.api import Api
|
||||
import uuid
|
||||
import asyncio
|
||||
from websockets.asyncio.client import connect
|
||||
|
||||
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
|
||||
default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
|
||||
|
||||
def query(url, flow_id, system, prompt):
|
||||
async def query(url, flow_id, system, prompt, streaming=True):
|
||||
|
||||
api = Api(url).flow().id(flow_id)
|
||||
if not url.endswith("/"):
|
||||
url += "/"
|
||||
|
||||
resp = api.text_completion(system=system, prompt=prompt)
|
||||
url = url + "api/v1/socket"
|
||||
|
||||
print(resp)
|
||||
mid = str(uuid.uuid4())
|
||||
|
||||
async with connect(url) as ws:
|
||||
|
||||
req = {
|
||||
"id": mid,
|
||||
"service": "text-completion",
|
||||
"flow": flow_id,
|
||||
"request": {
|
||||
"system": system,
|
||||
"prompt": prompt,
|
||||
"streaming": streaming
|
||||
}
|
||||
}
|
||||
|
||||
await ws.send(json.dumps(req))
|
||||
|
||||
while True:
|
||||
|
||||
msg = await ws.recv()
|
||||
|
||||
obj = json.loads(msg)
|
||||
|
||||
if "error" in obj:
|
||||
raise RuntimeError(obj["error"])
|
||||
|
||||
if obj["id"] != mid:
|
||||
continue
|
||||
|
||||
if "response" in obj["response"]:
|
||||
if streaming:
|
||||
# Stream output to stdout without newline
|
||||
print(obj["response"]["response"], end="", flush=True)
|
||||
else:
|
||||
# Non-streaming: print complete response
|
||||
print(obj["response"]["response"])
|
||||
|
||||
if obj["complete"]:
|
||||
if streaming:
|
||||
# Add final newline after streaming
|
||||
print()
|
||||
break
|
||||
|
||||
await ws.close()
|
||||
|
||||
def main():
|
||||
|
||||
|
|
@ -49,16 +95,23 @@ def main():
|
|||
help=f'Flow ID (default: default)'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--no-streaming',
|
||||
action='store_true',
|
||||
help='Disable streaming (default: streaming enabled)'
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
|
||||
query(
|
||||
asyncio.run(query(
|
||||
url=args.url,
|
||||
flow_id = args.flow_id,
|
||||
flow_id=args.flow_id,
|
||||
system=args.system[0],
|
||||
prompt=args.prompt[0],
|
||||
)
|
||||
streaming=not args.no_streaming
|
||||
))
|
||||
|
||||
except Exception as e:
|
||||
|
||||
|
|
|
|||
|
|
@ -10,20 +10,76 @@ using key=value arguments on the command line, and these replace
|
|||
import argparse
|
||||
import os
|
||||
import json
|
||||
from trustgraph.api import Api
|
||||
import uuid
|
||||
import asyncio
|
||||
from websockets.asyncio.client import connect
|
||||
|
||||
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
|
||||
default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
|
||||
|
||||
def query(url, flow_id, template_id, variables):
|
||||
async def query(url, flow_id, template_id, variables, streaming=True):
|
||||
|
||||
api = Api(url).flow().id(flow_id)
|
||||
if not url.endswith("/"):
|
||||
url += "/"
|
||||
|
||||
resp = api.prompt(id=template_id, variables=variables)
|
||||
url = url + "api/v1/socket"
|
||||
|
||||
if isinstance(resp, str):
|
||||
print(resp)
|
||||
else:
|
||||
print(json.dumps(resp, indent=4))
|
||||
mid = str(uuid.uuid4())
|
||||
|
||||
async with connect(url) as ws:
|
||||
|
||||
req = {
|
||||
"id": mid,
|
||||
"service": "prompt",
|
||||
"flow": flow_id,
|
||||
"request": {
|
||||
"id": template_id,
|
||||
"variables": variables,
|
||||
"streaming": streaming
|
||||
}
|
||||
}
|
||||
|
||||
await ws.send(json.dumps(req))
|
||||
|
||||
full_response = {"text": "", "object": ""}
|
||||
|
||||
while True:
|
||||
|
||||
msg = await ws.recv()
|
||||
|
||||
obj = json.loads(msg)
|
||||
|
||||
if "error" in obj:
|
||||
raise RuntimeError(obj["error"])
|
||||
|
||||
if obj["id"] != mid:
|
||||
continue
|
||||
|
||||
response = obj["response"]
|
||||
|
||||
# Handle text responses (streaming)
|
||||
if "text" in response and response["text"]:
|
||||
if streaming:
|
||||
# Stream output to stdout without newline
|
||||
print(response["text"], end="", flush=True)
|
||||
full_response["text"] += response["text"]
|
||||
else:
|
||||
# Non-streaming: print complete response
|
||||
print(response["text"])
|
||||
|
||||
# Handle object responses (JSON, never streamed)
|
||||
if "object" in response and response["object"]:
|
||||
full_response["object"] = response["object"]
|
||||
|
||||
if obj["complete"]:
|
||||
if streaming and full_response["text"]:
|
||||
# Add final newline after streaming text
|
||||
print()
|
||||
elif full_response["object"]:
|
||||
# Print JSON object (pretty-printed)
|
||||
print(json.dumps(json.loads(full_response["object"]), indent=4))
|
||||
break
|
||||
|
||||
await ws.close()
|
||||
|
||||
def main():
|
||||
|
||||
|
|
@ -59,6 +115,12 @@ def main():
|
|||
specified multiple times''',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--no-streaming',
|
||||
action='store_true',
|
||||
help='Disable streaming (default: streaming enabled for text responses)'
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
variables = {}
|
||||
|
|
@ -73,12 +135,13 @@ specified multiple times''',
|
|||
|
||||
try:
|
||||
|
||||
query(
|
||||
asyncio.run(query(
|
||||
url=args.url,
|
||||
flow_id=args.flow_id,
|
||||
template_id=args.id[0],
|
||||
variables=variables,
|
||||
)
|
||||
streaming=not args.no_streaming
|
||||
))
|
||||
|
||||
except Exception as e:
|
||||
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
import logging
|
||||
import json
|
||||
import re
|
||||
import asyncio
|
||||
|
||||
from . types import Action, Final
|
||||
|
||||
|
|
@ -169,7 +170,7 @@ class AgentManager:
|
|||
|
||||
raise ValueError(f"Could not parse response: {text}")
|
||||
|
||||
async def reason(self, question, history, context):
|
||||
async def reason(self, question, history, context, streaming=False, think=None, observe=None, answer=None):
|
||||
|
||||
logger.debug(f"calling reason: {question}")
|
||||
|
||||
|
|
@ -219,25 +220,62 @@ class AgentManager:
|
|||
|
||||
logger.info(f"prompt: {variables}")
|
||||
|
||||
# Get text response from prompt service
|
||||
response_text = await context("prompt-request").agent_react(variables)
|
||||
# Streaming path - use StreamingReActParser
|
||||
if streaming and think:
|
||||
from .streaming_parser import StreamingReActParser
|
||||
|
||||
logger.debug(f"Response text:\n{response_text}")
|
||||
# Create parser with streaming callbacks
|
||||
# Thought chunks go to think(), answer chunks go to answer()
|
||||
parser = StreamingReActParser(
|
||||
on_thought_chunk=lambda chunk: asyncio.create_task(think(chunk)),
|
||||
on_answer_chunk=lambda chunk: asyncio.create_task(answer(chunk) if answer else think(chunk)),
|
||||
)
|
||||
|
||||
logger.info(f"response: {response_text}")
|
||||
# Create async chunk callback that feeds parser
|
||||
async def on_chunk(text):
|
||||
parser.feed(text)
|
||||
|
||||
# Get streaming response
|
||||
response_text = await context("prompt-request").agent_react(
|
||||
variables=variables,
|
||||
streaming=True,
|
||||
chunk_callback=on_chunk
|
||||
)
|
||||
|
||||
# Finalize parser
|
||||
parser.finalize()
|
||||
|
||||
# Get result
|
||||
result = parser.get_result()
|
||||
if result is None:
|
||||
raise RuntimeError("Parser failed to produce a result")
|
||||
|
||||
# Parse the text response
|
||||
try:
|
||||
result = self.parse_react_response(response_text)
|
||||
logger.info(f"Parsed result: {result}")
|
||||
return result
|
||||
except ValueError as e:
|
||||
logger.error(f"Failed to parse response: {e}")
|
||||
# Try to provide a helpful error message
|
||||
logger.error(f"Response was: {response_text}")
|
||||
raise RuntimeError(f"Failed to parse agent response: {e}")
|
||||
|
||||
async def react(self, question, history, think, observe, context):
|
||||
else:
|
||||
# Non-streaming path - get complete text and parse
|
||||
response_text = await context("prompt-request").agent_react(
|
||||
variables=variables,
|
||||
streaming=False
|
||||
)
|
||||
|
||||
logger.debug(f"Response text:\n{response_text}")
|
||||
|
||||
logger.info(f"response: {response_text}")
|
||||
|
||||
# Parse the text response
|
||||
try:
|
||||
result = self.parse_react_response(response_text)
|
||||
logger.info(f"Parsed result: {result}")
|
||||
return result
|
||||
except ValueError as e:
|
||||
logger.error(f"Failed to parse response: {e}")
|
||||
# Try to provide a helpful error message
|
||||
logger.error(f"Response was: {response_text}")
|
||||
raise RuntimeError(f"Failed to parse agent response: {e}")
|
||||
|
||||
async def react(self, question, history, think, observe, context, streaming=False, answer=None):
|
||||
|
||||
logger.info(f"question: {question}")
|
||||
|
||||
|
|
@ -245,17 +283,27 @@ class AgentManager:
|
|||
question = question,
|
||||
history = history,
|
||||
context = context,
|
||||
streaming = streaming,
|
||||
think = think,
|
||||
observe = observe,
|
||||
answer = answer,
|
||||
)
|
||||
logger.info(f"act: {act}")
|
||||
|
||||
if isinstance(act, Final):
|
||||
|
||||
await think(act.thought)
|
||||
# In non-streaming mode, send complete thought
|
||||
# In streaming mode, thoughts were already sent as chunks
|
||||
if not streaming:
|
||||
await think(act.thought)
|
||||
return act
|
||||
|
||||
else:
|
||||
|
||||
await think(act.thought)
|
||||
# In non-streaming mode, send complete thought
|
||||
# In streaming mode, thoughts were already sent as chunks
|
||||
if not streaming:
|
||||
await think(act.thought)
|
||||
|
||||
logger.debug(f"ACTION: {act.name}")
|
||||
|
||||
|
|
|
|||
|
|
@ -191,6 +191,9 @@ class Processor(AgentService):
|
|||
|
||||
try:
|
||||
|
||||
# Check if streaming is enabled
|
||||
streaming = getattr(request, 'streaming', False)
|
||||
|
||||
if request.history:
|
||||
history = [
|
||||
Action(
|
||||
|
|
@ -215,12 +218,27 @@ class Processor(AgentService):
|
|||
|
||||
logger.debug(f"Think: {x}")
|
||||
|
||||
r = AgentResponse(
|
||||
answer=None,
|
||||
error=None,
|
||||
thought=x,
|
||||
observation=None,
|
||||
)
|
||||
if streaming:
|
||||
# Streaming format
|
||||
r = AgentResponse(
|
||||
chunk_type="thought",
|
||||
content=x,
|
||||
end_of_message=True,
|
||||
end_of_dialog=False,
|
||||
# Legacy fields for backward compatibility
|
||||
answer=None,
|
||||
error=None,
|
||||
thought=x,
|
||||
observation=None,
|
||||
)
|
||||
else:
|
||||
# Legacy format
|
||||
r = AgentResponse(
|
||||
answer=None,
|
||||
error=None,
|
||||
thought=x,
|
||||
observation=None,
|
||||
)
|
||||
|
||||
await respond(r)
|
||||
|
||||
|
|
@ -228,12 +246,55 @@ class Processor(AgentService):
|
|||
|
||||
logger.debug(f"Observe: {x}")
|
||||
|
||||
r = AgentResponse(
|
||||
answer=None,
|
||||
error=None,
|
||||
thought=None,
|
||||
observation=x,
|
||||
)
|
||||
if streaming:
|
||||
# Streaming format
|
||||
r = AgentResponse(
|
||||
chunk_type="observation",
|
||||
content=x,
|
||||
end_of_message=True,
|
||||
end_of_dialog=False,
|
||||
# Legacy fields for backward compatibility
|
||||
answer=None,
|
||||
error=None,
|
||||
thought=None,
|
||||
observation=x,
|
||||
)
|
||||
else:
|
||||
# Legacy format
|
||||
r = AgentResponse(
|
||||
answer=None,
|
||||
error=None,
|
||||
thought=None,
|
||||
observation=x,
|
||||
)
|
||||
|
||||
await respond(r)
|
||||
|
||||
async def answer(x):
|
||||
|
||||
logger.debug(f"Answer: {x}")
|
||||
|
||||
if streaming:
|
||||
# Streaming format
|
||||
r = AgentResponse(
|
||||
chunk_type="answer",
|
||||
content=x,
|
||||
end_of_message=False, # More chunks may follow
|
||||
end_of_dialog=False,
|
||||
# Legacy fields for backward compatibility
|
||||
answer=None,
|
||||
error=None,
|
||||
thought=None,
|
||||
observation=None,
|
||||
)
|
||||
else:
|
||||
# Legacy format - shouldn't be called in non-streaming mode
|
||||
r = AgentResponse(
|
||||
answer=x,
|
||||
error=None,
|
||||
thought=None,
|
||||
observation=None,
|
||||
)
|
||||
|
||||
await respond(r)
|
||||
|
||||
|
|
@ -273,7 +334,9 @@ class Processor(AgentService):
|
|||
history = history,
|
||||
think = think,
|
||||
observe = observe,
|
||||
answer = answer,
|
||||
context = UserAwareContext(flow, request.user),
|
||||
streaming = streaming,
|
||||
)
|
||||
|
||||
logger.debug(f"Action: {act}")
|
||||
|
|
@ -287,11 +350,26 @@ class Processor(AgentService):
|
|||
else:
|
||||
f = json.dumps(act.final)
|
||||
|
||||
r = AgentResponse(
|
||||
answer=act.final,
|
||||
error=None,
|
||||
thought=None,
|
||||
)
|
||||
if streaming:
|
||||
# Streaming format - send end-of-dialog marker
|
||||
# Answer chunks were already sent via think() callback during parsing
|
||||
r = AgentResponse(
|
||||
chunk_type="answer",
|
||||
content="", # Empty content, just marking end of dialog
|
||||
end_of_message=True,
|
||||
end_of_dialog=True,
|
||||
# Legacy fields for backward compatibility
|
||||
answer=act.final,
|
||||
error=None,
|
||||
thought=None,
|
||||
)
|
||||
else:
|
||||
# Legacy format - send complete answer
|
||||
r = AgentResponse(
|
||||
answer=act.final,
|
||||
error=None,
|
||||
thought=None,
|
||||
)
|
||||
|
||||
await respond(r)
|
||||
|
||||
|
|
@ -321,7 +399,9 @@ class Processor(AgentService):
|
|||
observation=h.observation
|
||||
)
|
||||
for h in history
|
||||
]
|
||||
],
|
||||
user=request.user,
|
||||
streaming=streaming,
|
||||
)
|
||||
|
||||
await next(r)
|
||||
|
|
@ -336,14 +416,32 @@ class Processor(AgentService):
|
|||
|
||||
logger.debug("Send error response...")
|
||||
|
||||
r = AgentResponse(
|
||||
error=Error(
|
||||
type = "agent-error",
|
||||
message = str(e),
|
||||
),
|
||||
response=None,
|
||||
error_obj = Error(
|
||||
type = "agent-error",
|
||||
message = str(e),
|
||||
)
|
||||
|
||||
# Check if streaming was enabled (may not be set if error occurred early)
|
||||
streaming = getattr(request, 'streaming', False) if 'request' in locals() else False
|
||||
|
||||
if streaming:
|
||||
# Streaming format
|
||||
r = AgentResponse(
|
||||
chunk_type="error",
|
||||
content=str(e),
|
||||
end_of_message=True,
|
||||
end_of_dialog=True,
|
||||
# Legacy fields for backward compatibility
|
||||
error=error_obj,
|
||||
response=None,
|
||||
)
|
||||
else:
|
||||
# Legacy format
|
||||
r = AgentResponse(
|
||||
error=error_obj,
|
||||
response=None,
|
||||
)
|
||||
|
||||
await respond(r)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
339
trustgraph-flow/trustgraph/agent/react/streaming_parser.py
Normal file
339
trustgraph-flow/trustgraph/agent/react/streaming_parser.py
Normal file
|
|
@ -0,0 +1,339 @@
|
|||
"""
|
||||
Streaming parser for ReAct responses.
|
||||
|
||||
This parser handles text chunks from LLM streaming responses and parses them
|
||||
into ReAct format (Thought/Action/Args or Thought/Final Answer). It maintains
|
||||
state across chunk boundaries to handle cases where delimiters or JSON are split.
|
||||
|
||||
Key challenges:
|
||||
- Delimiters may be split across chunks: "Tho" + "ught:" or "Final An" + "swer:"
|
||||
- JSON arguments may be split: '{"loc' + 'ation": "NYC"}'
|
||||
- Need to emit thought/answer chunks as they arrive for streaming
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from enum import Enum
|
||||
from typing import Optional, Callable, Any
|
||||
from . types import Action, Final
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ParserState(Enum):
|
||||
"""States for the streaming ReAct parser state machine"""
|
||||
INITIAL = "initial" # Waiting for first content
|
||||
THOUGHT = "thought" # Accumulating thought content
|
||||
ACTION = "action" # Found "Action:", collecting action name
|
||||
ARGS = "args" # Found "Args:", collecting JSON arguments
|
||||
FINAL_ANSWER = "final_answer" # Found "Final Answer:", collecting answer
|
||||
COMPLETE = "complete" # Parsing complete, object ready
|
||||
|
||||
|
||||
class StreamingReActParser:
|
||||
"""
|
||||
Stateful parser for streaming ReAct responses.
|
||||
|
||||
Expected format:
|
||||
Thought: [reasoning about what to do next]
|
||||
Action: [tool_name]
|
||||
Args: {
|
||||
"param": "value"
|
||||
}
|
||||
|
||||
OR
|
||||
Thought: [reasoning about the final answer]
|
||||
Final Answer: [the answer]
|
||||
|
||||
Usage:
|
||||
parser = StreamingReActParser(
|
||||
on_thought_chunk=lambda chunk: print(f"Thought: {chunk}"),
|
||||
on_answer_chunk=lambda chunk: print(f"Answer: {chunk}"),
|
||||
)
|
||||
|
||||
for chunk in llm_stream:
|
||||
parser.feed(chunk)
|
||||
if parser.is_complete():
|
||||
result = parser.get_result()
|
||||
break
|
||||
"""
|
||||
|
||||
# Delimiters we're looking for
|
||||
THOUGHT_DELIMITER = "Thought:"
|
||||
ACTION_DELIMITER = "Action:"
|
||||
ARGS_DELIMITER = "Args:"
|
||||
FINAL_ANSWER_DELIMITER = "Final Answer:"
|
||||
|
||||
# Maximum buffer size for delimiter detection (longest delimiter + safety margin)
|
||||
MAX_DELIMITER_BUFFER = 20
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
on_thought_chunk: Optional[Callable[[str], Any]] = None,
|
||||
on_answer_chunk: Optional[Callable[[str], Any]] = None,
|
||||
):
|
||||
"""
|
||||
Initialize streaming parser.
|
||||
|
||||
Args:
|
||||
on_thought_chunk: Callback for thought text chunks as they arrive
|
||||
on_answer_chunk: Callback for final answer text chunks as they arrive
|
||||
"""
|
||||
self.on_thought_chunk = on_thought_chunk
|
||||
self.on_answer_chunk = on_answer_chunk
|
||||
|
||||
# Parser state
|
||||
self.state = ParserState.INITIAL
|
||||
|
||||
# Buffers for accumulating content
|
||||
self.line_buffer = "" # For detecting delimiters across chunk boundaries
|
||||
self.thought_buffer = "" # Accumulated thought text
|
||||
self.action_buffer = "" # Action name
|
||||
self.args_buffer = "" # JSON arguments text
|
||||
self.answer_buffer = "" # Final answer text
|
||||
|
||||
# JSON parsing state for Args
|
||||
self.brace_count = 0
|
||||
self.args_started = False
|
||||
|
||||
# Result object (Action or Final)
|
||||
self.result = None
|
||||
|
||||
def feed(self, chunk: str) -> None:
|
||||
"""
|
||||
Feed a text chunk to the parser.
|
||||
|
||||
Args:
|
||||
chunk: Text chunk from LLM stream
|
||||
"""
|
||||
if self.state == ParserState.COMPLETE:
|
||||
return # Already complete, ignore further chunks
|
||||
|
||||
# Add chunk to line buffer for delimiter detection
|
||||
self.line_buffer += chunk
|
||||
|
||||
# Remove markdown code blocks if present
|
||||
self.line_buffer = re.sub(r'^```[^\n]*\n', '', self.line_buffer)
|
||||
self.line_buffer = re.sub(r'\n```$', '', self.line_buffer)
|
||||
|
||||
# Process based on current state
|
||||
while self.line_buffer and self.state != ParserState.COMPLETE:
|
||||
if self.state == ParserState.INITIAL:
|
||||
self._process_initial()
|
||||
elif self.state == ParserState.THOUGHT:
|
||||
self._process_thought()
|
||||
elif self.state == ParserState.ACTION:
|
||||
self._process_action()
|
||||
elif self.state == ParserState.ARGS:
|
||||
self._process_args()
|
||||
elif self.state == ParserState.FINAL_ANSWER:
|
||||
self._process_final_answer()
|
||||
|
||||
def _process_initial(self) -> None:
|
||||
"""Process INITIAL state - looking for 'Thought:' delimiter"""
|
||||
idx = self.line_buffer.find(self.THOUGHT_DELIMITER)
|
||||
|
||||
if idx >= 0:
|
||||
# Found thought delimiter
|
||||
# Discard any content before it
|
||||
self.line_buffer = self.line_buffer[idx + len(self.THOUGHT_DELIMITER):]
|
||||
self.state = ParserState.THOUGHT
|
||||
elif len(self.line_buffer) >= self.MAX_DELIMITER_BUFFER:
|
||||
# Buffer getting too large, probably junk before thought
|
||||
# Keep only the tail that might contain partial delimiter
|
||||
self.line_buffer = self.line_buffer[-self.MAX_DELIMITER_BUFFER:]
|
||||
|
||||
def _process_thought(self) -> None:
|
||||
"""Process THOUGHT state - accumulating thought content"""
|
||||
# Check for Action or Final Answer delimiter
|
||||
action_idx = self.line_buffer.find(self.ACTION_DELIMITER)
|
||||
final_idx = self.line_buffer.find(self.FINAL_ANSWER_DELIMITER)
|
||||
|
||||
# Find which delimiter comes first (if any)
|
||||
next_delimiter_idx = -1
|
||||
next_state = None
|
||||
|
||||
if action_idx >= 0 and (final_idx < 0 or action_idx < final_idx):
|
||||
next_delimiter_idx = action_idx
|
||||
next_state = ParserState.ACTION
|
||||
delimiter_len = len(self.ACTION_DELIMITER)
|
||||
elif final_idx >= 0:
|
||||
next_delimiter_idx = final_idx
|
||||
next_state = ParserState.FINAL_ANSWER
|
||||
delimiter_len = len(self.FINAL_ANSWER_DELIMITER)
|
||||
|
||||
if next_delimiter_idx >= 0:
|
||||
# Found next delimiter
|
||||
thought_chunk = self.line_buffer[:next_delimiter_idx].strip()
|
||||
if thought_chunk:
|
||||
self.thought_buffer += thought_chunk
|
||||
if self.on_thought_chunk:
|
||||
self.on_thought_chunk(thought_chunk)
|
||||
|
||||
self.line_buffer = self.line_buffer[next_delimiter_idx + delimiter_len:]
|
||||
self.state = next_state
|
||||
else:
|
||||
# No delimiter found yet
|
||||
# Keep tail in buffer (might contain partial delimiter)
|
||||
# Emit the rest as thought chunk
|
||||
if len(self.line_buffer) > self.MAX_DELIMITER_BUFFER:
|
||||
emittable = self.line_buffer[:-self.MAX_DELIMITER_BUFFER]
|
||||
self.thought_buffer += emittable
|
||||
if self.on_thought_chunk:
|
||||
self.on_thought_chunk(emittable)
|
||||
self.line_buffer = self.line_buffer[-self.MAX_DELIMITER_BUFFER:]
|
||||
|
||||
def _process_action(self) -> None:
|
||||
"""Process ACTION state - collecting action name"""
|
||||
# Action name is on one line (or at least until newline or Args:)
|
||||
newline_idx = self.line_buffer.find('\n')
|
||||
args_idx = self.line_buffer.find(self.ARGS_DELIMITER)
|
||||
|
||||
# Find which comes first
|
||||
if args_idx >= 0 and (newline_idx < 0 or args_idx < newline_idx):
|
||||
# Args delimiter found first
|
||||
self.action_buffer = self.line_buffer[:args_idx].strip().strip('"')
|
||||
self.line_buffer = self.line_buffer[args_idx + len(self.ARGS_DELIMITER):]
|
||||
self.state = ParserState.ARGS
|
||||
elif newline_idx >= 0:
|
||||
# Newline found, action name complete
|
||||
self.action_buffer = self.line_buffer[:newline_idx].strip().strip('"')
|
||||
self.line_buffer = self.line_buffer[newline_idx + 1:]
|
||||
# Stay in ACTION state or move to ARGS if we find delimiter
|
||||
# Actually, check if next line has Args:
|
||||
if self.line_buffer.lstrip().startswith(self.ARGS_DELIMITER):
|
||||
args_start = self.line_buffer.find(self.ARGS_DELIMITER)
|
||||
self.line_buffer = self.line_buffer[args_start + len(self.ARGS_DELIMITER):]
|
||||
self.state = ParserState.ARGS
|
||||
else:
|
||||
# Not enough content yet, keep buffering
|
||||
# But if buffer is getting large, action name is probably complete
|
||||
if len(self.line_buffer) > 100:
|
||||
self.action_buffer = self.line_buffer.strip().strip('"')
|
||||
self.line_buffer = ""
|
||||
# Assume Args comes next, but we need more content
|
||||
self.state = ParserState.ARGS
|
||||
|
||||
def _process_args(self) -> None:
|
||||
"""Process ARGS state - collecting JSON arguments"""
|
||||
# Process character by character to track brace matching
|
||||
i = 0
|
||||
while i < len(self.line_buffer):
|
||||
char = self.line_buffer[i]
|
||||
self.args_buffer += char
|
||||
|
||||
if char == '{':
|
||||
self.brace_count += 1
|
||||
self.args_started = True
|
||||
elif char == '}':
|
||||
self.brace_count -= 1
|
||||
|
||||
# Check if JSON is complete
|
||||
if self.args_started and self.brace_count == 0:
|
||||
# JSON complete, try to parse
|
||||
try:
|
||||
args_dict = json.loads(self.args_buffer.strip())
|
||||
# Success! Create Action result
|
||||
self.result = Action(
|
||||
thought=self.thought_buffer.strip(),
|
||||
name=self.action_buffer,
|
||||
arguments=args_dict,
|
||||
observation=""
|
||||
)
|
||||
self.state = ParserState.COMPLETE
|
||||
self.line_buffer = "" # Clear buffer
|
||||
return
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"Failed to parse JSON args: {self.args_buffer}")
|
||||
raise ValueError(f"Invalid JSON in Args: {e}")
|
||||
|
||||
i += 1
|
||||
|
||||
# Consumed entire buffer, clear it and wait for more chunks
|
||||
self.line_buffer = ""
|
||||
|
||||
def _process_final_answer(self) -> None:
|
||||
"""Process FINAL_ANSWER state - collecting final answer"""
|
||||
# For final answer, we consume everything until we decide we're done
|
||||
# In streaming mode, we can't know when answer is complete until stream ends
|
||||
# So we emit chunks and accumulate
|
||||
|
||||
# Check if this might be JSON
|
||||
is_json = self.answer_buffer.strip().startswith('{') or \
|
||||
self.line_buffer.strip().startswith('{')
|
||||
|
||||
if is_json:
|
||||
# Handle JSON final answer
|
||||
self.answer_buffer += self.line_buffer
|
||||
|
||||
# Count braces to detect completion
|
||||
brace_count = self.answer_buffer.count('{') - self.answer_buffer.count('}')
|
||||
|
||||
if brace_count == 0 and '{' in self.answer_buffer:
|
||||
# JSON might be complete
|
||||
# Note: We can't be 100% sure without trying to parse
|
||||
# But in streaming mode, we'll finish when stream ends
|
||||
pass
|
||||
|
||||
# Emit chunk
|
||||
if self.on_answer_chunk:
|
||||
self.on_answer_chunk(self.line_buffer)
|
||||
|
||||
self.line_buffer = ""
|
||||
else:
|
||||
# Regular text answer - emit everything
|
||||
if self.line_buffer:
|
||||
self.answer_buffer += self.line_buffer
|
||||
if self.on_answer_chunk:
|
||||
self.on_answer_chunk(self.line_buffer)
|
||||
self.line_buffer = ""
|
||||
|
||||
def finalize(self) -> None:
|
||||
"""
|
||||
Call this when the stream is complete to finalize parsing.
|
||||
This handles any remaining buffered content.
|
||||
"""
|
||||
if self.state == ParserState.COMPLETE:
|
||||
return
|
||||
|
||||
# Flush any remaining thought chunks
|
||||
if self.state == ParserState.THOUGHT and self.line_buffer:
|
||||
self.thought_buffer += self.line_buffer
|
||||
if self.on_thought_chunk:
|
||||
self.on_thought_chunk(self.line_buffer)
|
||||
self.line_buffer = ""
|
||||
|
||||
# Finalize final answer
|
||||
if self.state == ParserState.FINAL_ANSWER:
|
||||
# Flush any remaining answer content
|
||||
if self.line_buffer:
|
||||
self.answer_buffer += self.line_buffer
|
||||
if self.on_answer_chunk:
|
||||
self.on_answer_chunk(self.line_buffer)
|
||||
self.line_buffer = ""
|
||||
|
||||
# Create Final result
|
||||
self.result = Final(
|
||||
thought=self.thought_buffer.strip(),
|
||||
final=self.answer_buffer.strip()
|
||||
)
|
||||
self.state = ParserState.COMPLETE
|
||||
|
||||
# If we're in other states, something went wrong
|
||||
if self.state not in [ParserState.COMPLETE, ParserState.FINAL_ANSWER]:
|
||||
if self.thought_buffer:
|
||||
raise ValueError(
|
||||
f"Stream ended in {self.state.value} state with incomplete parsing. "
|
||||
f"Thought: {self.thought_buffer[:100]}..."
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Stream ended in {self.state.value} state with no content")
|
||||
|
||||
def is_complete(self) -> bool:
|
||||
"""Check if parsing is complete"""
|
||||
return self.state == ParserState.COMPLETE
|
||||
|
||||
def get_result(self) -> Optional[Action | Final]:
|
||||
"""Get the parsed result (Action or Final)"""
|
||||
return self.result
|
||||
|
|
@ -34,9 +34,9 @@ class BlobStore:
|
|||
def ensure_bucket(self):
|
||||
|
||||
# Make the bucket if it doesn't exist.
|
||||
found = self.minio.bucket_exists(self.bucket_name)
|
||||
found = self.minio.bucket_exists(bucket_name=self.bucket_name)
|
||||
if not found:
|
||||
self.minio.make_bucket(self.bucket_name)
|
||||
self.minio.make_bucket(bucket_name=self.bucket_name)
|
||||
logger.info(f"Created bucket {self.bucket_name}")
|
||||
else:
|
||||
logger.debug(f"Bucket {self.bucket_name} already exists")
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ import os
|
|||
import logging
|
||||
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import LlmService, LlmResult
|
||||
from .... base import LlmService, LlmResult, LlmChunk
|
||||
|
||||
# Module logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -55,7 +55,7 @@ class Processor(LlmService):
|
|||
self.max_output = max_output
|
||||
self.default_model = model
|
||||
|
||||
def build_prompt(self, system, content, temperature=None):
|
||||
def build_prompt(self, system, content, temperature=None, stream=False):
|
||||
# Use provided temperature or fall back to default
|
||||
effective_temperature = temperature if temperature is not None else self.temperature
|
||||
|
||||
|
|
@ -73,6 +73,9 @@ class Processor(LlmService):
|
|||
"top_p": 1
|
||||
}
|
||||
|
||||
if stream:
|
||||
data["stream"] = True
|
||||
|
||||
body = json.dumps(data)
|
||||
|
||||
return body
|
||||
|
|
@ -157,6 +160,84 @@ class Processor(LlmService):
|
|||
|
||||
logger.debug("Azure LLM processing complete")
|
||||
|
||||
def supports_streaming(self):
|
||||
"""Azure serverless endpoints support streaming"""
|
||||
return True
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
||||
"""Stream content generation from Azure serverless endpoint"""
|
||||
model_name = model or self.default_model
|
||||
effective_temperature = temperature if temperature is not None else self.temperature
|
||||
|
||||
logger.debug(f"Using model (streaming): {model_name}")
|
||||
logger.debug(f"Using temperature: {effective_temperature}")
|
||||
|
||||
try:
|
||||
body = self.build_prompt(system, prompt, effective_temperature, stream=True)
|
||||
|
||||
url = self.endpoint
|
||||
api_key = self.token
|
||||
|
||||
headers = {
|
||||
'Content-Type': 'application/json',
|
||||
'Authorization': f'Bearer {api_key}'
|
||||
}
|
||||
|
||||
response = requests.post(url, data=body, headers=headers, stream=True)
|
||||
|
||||
if response.status_code == 429:
|
||||
raise TooManyRequests()
|
||||
|
||||
if response.status_code != 200:
|
||||
raise RuntimeError("LLM failure")
|
||||
|
||||
# Parse SSE stream
|
||||
for line in response.iter_lines():
|
||||
if line:
|
||||
line = line.decode('utf-8').strip()
|
||||
if line.startswith('data: '):
|
||||
data = line[6:] # Remove 'data: ' prefix
|
||||
|
||||
if data == '[DONE]':
|
||||
break
|
||||
|
||||
try:
|
||||
chunk_data = json.loads(data)
|
||||
|
||||
if 'choices' in chunk_data and len(chunk_data['choices']) > 0:
|
||||
delta = chunk_data['choices'][0].get('delta', {})
|
||||
content = delta.get('content')
|
||||
if content:
|
||||
yield LlmChunk(
|
||||
text=content,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=False
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Failed to parse chunk: {data}")
|
||||
continue
|
||||
|
||||
# Send final chunk
|
||||
yield LlmChunk(
|
||||
text="",
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=True
|
||||
)
|
||||
|
||||
logger.debug("Streaming complete")
|
||||
|
||||
except TooManyRequests:
|
||||
logger.warning("Rate limit exceeded during streaming")
|
||||
raise TooManyRequests()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Azure streaming exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ import logging
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import LlmService, LlmResult
|
||||
from .... base import LlmService, LlmResult, LlmChunk
|
||||
|
||||
default_ident = "text-completion"
|
||||
|
||||
|
|
@ -125,6 +125,75 @@ class Processor(LlmService):
|
|||
|
||||
logger.debug("Azure OpenAI LLM processing complete")
|
||||
|
||||
def supports_streaming(self):
|
||||
"""Azure OpenAI supports streaming"""
|
||||
return True
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
||||
"""
|
||||
Stream content generation from Azure OpenAI.
|
||||
Yields LlmChunk objects with is_final=True on the last chunk.
|
||||
"""
|
||||
# Use provided model or fall back to default
|
||||
model_name = model or self.default_model
|
||||
# Use provided temperature or fall back to default
|
||||
effective_temperature = temperature if temperature is not None else self.temperature
|
||||
|
||||
logger.debug(f"Using model (streaming): {model_name}")
|
||||
logger.debug(f"Using temperature: {effective_temperature}")
|
||||
|
||||
prompt = system + "\n\n" + prompt
|
||||
|
||||
try:
|
||||
response = self.openai.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": prompt
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
temperature=effective_temperature,
|
||||
max_tokens=self.max_output,
|
||||
top_p=1,
|
||||
stream=True # Enable streaming
|
||||
)
|
||||
|
||||
# Stream chunks
|
||||
for chunk in response:
|
||||
if chunk.choices and chunk.choices[0].delta.content:
|
||||
yield LlmChunk(
|
||||
text=chunk.choices[0].delta.content,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=False
|
||||
)
|
||||
|
||||
# Send final chunk
|
||||
yield LlmChunk(
|
||||
text="",
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=True
|
||||
)
|
||||
|
||||
logger.debug("Streaming complete")
|
||||
|
||||
except RateLimitError:
|
||||
logger.warning("Rate limit exceeded during streaming")
|
||||
raise TooManyRequests()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Azure OpenAI streaming exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ import os
|
|||
import logging
|
||||
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import LlmService, LlmResult
|
||||
from .... base import LlmService, LlmResult, LlmChunk
|
||||
|
||||
# Module logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -106,6 +106,65 @@ class Processor(LlmService):
|
|||
logger.error(f"Claude LLM exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
def supports_streaming(self):
|
||||
"""Claude/Anthropic supports streaming"""
|
||||
return True
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
||||
"""Stream content generation from Claude"""
|
||||
model_name = model or self.default_model
|
||||
effective_temperature = temperature if temperature is not None else self.temperature
|
||||
|
||||
logger.debug(f"Using model (streaming): {model_name}")
|
||||
logger.debug(f"Using temperature: {effective_temperature}")
|
||||
|
||||
try:
|
||||
with self.claude.messages.stream(
|
||||
model=model_name,
|
||||
max_tokens=self.max_output,
|
||||
temperature=effective_temperature,
|
||||
system=system,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": prompt
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
) as stream:
|
||||
for text in stream.text_stream:
|
||||
yield LlmChunk(
|
||||
text=text,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=False
|
||||
)
|
||||
|
||||
# Get final message for token counts
|
||||
final_message = stream.get_final_message()
|
||||
yield LlmChunk(
|
||||
text="",
|
||||
in_token=final_message.usage.input_tokens,
|
||||
out_token=final_message.usage.output_tokens,
|
||||
model=model_name,
|
||||
is_final=True
|
||||
)
|
||||
|
||||
logger.debug("Streaming complete")
|
||||
|
||||
except anthropic.RateLimitError:
|
||||
logger.warning("Rate limit exceeded during streaming")
|
||||
raise TooManyRequests()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Claude streaming exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ import logging
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import LlmService, LlmResult
|
||||
from .... base import LlmService, LlmResult, LlmChunk
|
||||
|
||||
default_ident = "text-completion"
|
||||
|
||||
|
|
@ -98,6 +98,68 @@ class Processor(LlmService):
|
|||
logger.error(f"Cohere LLM exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
def supports_streaming(self):
|
||||
"""Cohere supports streaming"""
|
||||
return True
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
||||
"""Stream content generation from Cohere"""
|
||||
model_name = model or self.default_model
|
||||
effective_temperature = temperature if temperature is not None else self.temperature
|
||||
|
||||
logger.debug(f"Using model (streaming): {model_name}")
|
||||
logger.debug(f"Using temperature: {effective_temperature}")
|
||||
|
||||
try:
|
||||
stream = self.cohere.chat_stream(
|
||||
model=model_name,
|
||||
message=prompt,
|
||||
preamble=system,
|
||||
temperature=effective_temperature,
|
||||
chat_history=[],
|
||||
prompt_truncation='auto',
|
||||
connectors=[]
|
||||
)
|
||||
|
||||
total_input_tokens = 0
|
||||
total_output_tokens = 0
|
||||
|
||||
for event in stream:
|
||||
if event.event_type == "text-generation":
|
||||
if hasattr(event, 'text') and event.text:
|
||||
yield LlmChunk(
|
||||
text=event.text,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=False
|
||||
)
|
||||
elif event.event_type == "stream-end":
|
||||
# Extract token counts from final event
|
||||
if hasattr(event, 'response') and hasattr(event.response, 'meta'):
|
||||
if hasattr(event.response.meta, 'billed_units'):
|
||||
total_input_tokens = int(event.response.meta.billed_units.input_tokens)
|
||||
total_output_tokens = int(event.response.meta.billed_units.output_tokens)
|
||||
|
||||
# Send final chunk with token counts
|
||||
yield LlmChunk(
|
||||
text="",
|
||||
in_token=total_input_tokens,
|
||||
out_token=total_output_tokens,
|
||||
model=model_name,
|
||||
is_final=True
|
||||
)
|
||||
|
||||
logger.debug("Streaming complete")
|
||||
|
||||
except cohere.TooManyRequestsError:
|
||||
logger.warning("Rate limit exceeded during streaming")
|
||||
raise TooManyRequests()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Cohere streaming exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ import logging
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import LlmService, LlmResult
|
||||
from .... base import LlmService, LlmResult, LlmChunk
|
||||
|
||||
default_ident = "text-completion"
|
||||
|
||||
|
|
@ -159,6 +159,67 @@ class Processor(LlmService):
|
|||
logger.error(f"GoogleAIStudio LLM exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
def supports_streaming(self):
|
||||
"""Google AI Studio supports streaming"""
|
||||
return True
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
||||
"""Stream content generation from Google AI Studio"""
|
||||
model_name = model or self.default_model
|
||||
effective_temperature = temperature if temperature is not None else self.temperature
|
||||
|
||||
logger.debug(f"Using model (streaming): {model_name}")
|
||||
logger.debug(f"Using temperature: {effective_temperature}")
|
||||
|
||||
generation_config = self._get_or_create_config(model_name, effective_temperature)
|
||||
generation_config.system_instruction = system
|
||||
|
||||
try:
|
||||
response = self.client.models.generate_content_stream(
|
||||
model=model_name,
|
||||
config=generation_config,
|
||||
contents=prompt,
|
||||
)
|
||||
|
||||
total_input_tokens = 0
|
||||
total_output_tokens = 0
|
||||
|
||||
for chunk in response:
|
||||
if hasattr(chunk, 'text') and chunk.text:
|
||||
yield LlmChunk(
|
||||
text=chunk.text,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=False
|
||||
)
|
||||
|
||||
# Accumulate token counts if available
|
||||
if hasattr(chunk, 'usage_metadata'):
|
||||
if hasattr(chunk.usage_metadata, 'prompt_token_count'):
|
||||
total_input_tokens = int(chunk.usage_metadata.prompt_token_count)
|
||||
if hasattr(chunk.usage_metadata, 'candidates_token_count'):
|
||||
total_output_tokens = int(chunk.usage_metadata.candidates_token_count)
|
||||
|
||||
# Send final chunk with token counts
|
||||
yield LlmChunk(
|
||||
text="",
|
||||
in_token=total_input_tokens,
|
||||
out_token=total_output_tokens,
|
||||
model=model_name,
|
||||
is_final=True
|
||||
)
|
||||
|
||||
logger.debug("Streaming complete")
|
||||
|
||||
except ResourceExhausted:
|
||||
logger.warning("Rate limit exceeded during streaming")
|
||||
raise TooManyRequests()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"GoogleAIStudio streaming exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import logging
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import LlmService, LlmResult
|
||||
from .... base import LlmService, LlmResult, LlmChunk
|
||||
|
||||
default_ident = "text-completion"
|
||||
|
||||
|
|
@ -102,6 +102,57 @@ class Processor(LlmService):
|
|||
logger.error(f"Llamafile LLM exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
def supports_streaming(self):
|
||||
"""LlamaFile supports streaming"""
|
||||
return True
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
||||
"""Stream content generation from LlamaFile"""
|
||||
model_name = model or self.default_model
|
||||
effective_temperature = temperature if temperature is not None else self.temperature
|
||||
|
||||
logger.debug(f"Using model (streaming): {model_name}")
|
||||
logger.debug(f"Using temperature: {effective_temperature}")
|
||||
|
||||
prompt = system + "\n\n" + prompt
|
||||
|
||||
try:
|
||||
response = self.openai.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
temperature=effective_temperature,
|
||||
max_tokens=self.max_output,
|
||||
top_p=1,
|
||||
frequency_penalty=0,
|
||||
presence_penalty=0,
|
||||
response_format={"type": "text"},
|
||||
stream=True
|
||||
)
|
||||
|
||||
for chunk in response:
|
||||
if chunk.choices and chunk.choices[0].delta.content:
|
||||
yield LlmChunk(
|
||||
text=chunk.choices[0].delta.content,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=False
|
||||
)
|
||||
|
||||
yield LlmChunk(
|
||||
text="",
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=True
|
||||
)
|
||||
|
||||
logger.debug("Streaming complete")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"LlamaFile streaming exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import logging
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import LlmService, LlmResult
|
||||
from .... base import LlmService, LlmResult, LlmChunk
|
||||
|
||||
default_ident = "text-completion"
|
||||
|
||||
|
|
@ -106,6 +106,57 @@ class Processor(LlmService):
|
|||
logger.error(f"LMStudio LLM exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
def supports_streaming(self):
|
||||
"""LM Studio supports streaming"""
|
||||
return True
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
||||
"""Stream content generation from LM Studio"""
|
||||
model_name = model or self.default_model
|
||||
effective_temperature = temperature if temperature is not None else self.temperature
|
||||
|
||||
logger.debug(f"Using model (streaming): {model_name}")
|
||||
logger.debug(f"Using temperature: {effective_temperature}")
|
||||
|
||||
prompt = system + "\n\n" + prompt
|
||||
|
||||
try:
|
||||
response = self.openai.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
temperature=effective_temperature,
|
||||
max_tokens=self.max_output,
|
||||
top_p=1,
|
||||
frequency_penalty=0,
|
||||
presence_penalty=0,
|
||||
response_format={"type": "text"},
|
||||
stream=True
|
||||
)
|
||||
|
||||
for chunk in response:
|
||||
if chunk.choices and chunk.choices[0].delta.content:
|
||||
yield LlmChunk(
|
||||
text=chunk.choices[0].delta.content,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=False
|
||||
)
|
||||
|
||||
yield LlmChunk(
|
||||
text="",
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=True
|
||||
)
|
||||
|
||||
logger.debug("Streaming complete")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"LMStudio streaming exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import logging
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import LlmService, LlmResult
|
||||
from .... base import LlmService, LlmResult, LlmChunk
|
||||
|
||||
default_ident = "text-completion"
|
||||
|
||||
|
|
@ -120,6 +120,67 @@ class Processor(LlmService):
|
|||
logger.error(f"Mistral LLM exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
def supports_streaming(self):
|
||||
"""Mistral supports streaming"""
|
||||
return True
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
||||
"""Stream content generation from Mistral"""
|
||||
model_name = model or self.default_model
|
||||
effective_temperature = temperature if temperature is not None else self.temperature
|
||||
|
||||
logger.debug(f"Using model (streaming): {model_name}")
|
||||
logger.debug(f"Using temperature: {effective_temperature}")
|
||||
|
||||
prompt = system + "\n\n" + prompt
|
||||
|
||||
try:
|
||||
stream = self.mistral.chat.stream(
|
||||
model=model_name,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": prompt
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
temperature=effective_temperature,
|
||||
max_tokens=self.max_output,
|
||||
top_p=1,
|
||||
frequency_penalty=0,
|
||||
presence_penalty=0,
|
||||
response_format={"type": "text"}
|
||||
)
|
||||
|
||||
for chunk in stream:
|
||||
if chunk.data.choices and chunk.data.choices[0].delta.content:
|
||||
yield LlmChunk(
|
||||
text=chunk.data.choices[0].delta.content,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=False
|
||||
)
|
||||
|
||||
# Send final chunk
|
||||
yield LlmChunk(
|
||||
text="",
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=True
|
||||
)
|
||||
|
||||
logger.debug("Streaming complete")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Mistral streaming exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import logging
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import LlmService, LlmResult
|
||||
from .... base import LlmService, LlmResult, LlmChunk
|
||||
|
||||
default_ident = "text-completion"
|
||||
|
||||
|
|
@ -79,6 +79,62 @@ class Processor(LlmService):
|
|||
logger.error(f"Ollama LLM exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
def supports_streaming(self):
|
||||
"""Ollama supports streaming"""
|
||||
return True
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
||||
"""Stream content generation from Ollama"""
|
||||
model_name = model or self.default_model
|
||||
effective_temperature = temperature if temperature is not None else self.temperature
|
||||
|
||||
logger.debug(f"Using model (streaming): {model_name}")
|
||||
logger.debug(f"Using temperature: {effective_temperature}")
|
||||
|
||||
prompt = system + "\n\n" + prompt
|
||||
|
||||
try:
|
||||
stream = self.llm.generate(
|
||||
model_name,
|
||||
prompt,
|
||||
options={'temperature': effective_temperature},
|
||||
stream=True
|
||||
)
|
||||
|
||||
total_input_tokens = 0
|
||||
total_output_tokens = 0
|
||||
|
||||
for chunk in stream:
|
||||
if 'response' in chunk and chunk['response']:
|
||||
yield LlmChunk(
|
||||
text=chunk['response'],
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=False
|
||||
)
|
||||
|
||||
# Accumulate token counts if available
|
||||
if 'prompt_eval_count' in chunk:
|
||||
total_input_tokens = int(chunk['prompt_eval_count'])
|
||||
if 'eval_count' in chunk:
|
||||
total_output_tokens = int(chunk['eval_count'])
|
||||
|
||||
# Send final chunk with token counts
|
||||
yield LlmChunk(
|
||||
text="",
|
||||
in_token=total_input_tokens,
|
||||
out_token=total_output_tokens,
|
||||
model=model_name,
|
||||
is_final=True
|
||||
)
|
||||
|
||||
logger.debug("Streaming complete")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Ollama streaming exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ import os
|
|||
import logging
|
||||
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import LlmService, LlmResult
|
||||
from .... base import LlmService, LlmResult, LlmChunk
|
||||
|
||||
# Module logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -118,6 +118,75 @@ class Processor(LlmService):
|
|||
logger.error(f"OpenAI LLM exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
def supports_streaming(self):
|
||||
"""OpenAI supports streaming"""
|
||||
return True
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
||||
"""
|
||||
Stream content generation from OpenAI.
|
||||
Yields LlmChunk objects with is_final=True on the last chunk.
|
||||
"""
|
||||
# Use provided model or fall back to default
|
||||
model_name = model or self.default_model
|
||||
# Use provided temperature or fall back to default
|
||||
effective_temperature = temperature if temperature is not None else self.temperature
|
||||
|
||||
logger.debug(f"Using model (streaming): {model_name}")
|
||||
logger.debug(f"Using temperature: {effective_temperature}")
|
||||
|
||||
prompt = system + "\n\n" + prompt
|
||||
|
||||
try:
|
||||
response = self.openai.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": prompt
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
temperature=effective_temperature,
|
||||
max_tokens=self.max_output,
|
||||
stream=True # Enable streaming
|
||||
)
|
||||
|
||||
# Stream chunks
|
||||
for chunk in response:
|
||||
if chunk.choices and chunk.choices[0].delta.content:
|
||||
yield LlmChunk(
|
||||
text=chunk.choices[0].delta.content,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=False
|
||||
)
|
||||
|
||||
# Note: OpenAI doesn't provide token counts in streaming mode
|
||||
# Send final chunk without token counts
|
||||
yield LlmChunk(
|
||||
text="",
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=True
|
||||
)
|
||||
|
||||
logger.debug("Streaming complete")
|
||||
|
||||
except RateLimitError:
|
||||
logger.warning("Hit rate limit during streaming")
|
||||
raise TooManyRequests()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"OpenAI streaming exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import logging
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import LlmService, LlmResult
|
||||
from .... base import LlmService, LlmResult, LlmChunk
|
||||
|
||||
default_ident = "text-completion"
|
||||
|
||||
|
|
@ -121,6 +121,100 @@ class Processor(LlmService):
|
|||
logger.error(f"TGI LLM exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
def supports_streaming(self):
|
||||
"""TGI supports streaming"""
|
||||
return True
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
||||
"""Stream content generation from TGI"""
|
||||
model_name = model or self.default_model
|
||||
effective_temperature = temperature if temperature is not None else self.temperature
|
||||
|
||||
logger.debug(f"Using model (streaming): {model_name}")
|
||||
logger.debug(f"Using temperature: {effective_temperature}")
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
request = {
|
||||
"model": model_name,
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": system,
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt,
|
||||
}
|
||||
],
|
||||
"max_tokens": self.max_output,
|
||||
"temperature": effective_temperature,
|
||||
"stream": True,
|
||||
}
|
||||
|
||||
try:
|
||||
url = f"{self.base_url}/chat/completions"
|
||||
|
||||
async with self.session.post(
|
||||
url,
|
||||
headers=headers,
|
||||
json=request,
|
||||
) as response:
|
||||
|
||||
if response.status != 200:
|
||||
raise RuntimeError("Bad status: " + str(response.status))
|
||||
|
||||
# Parse SSE stream
|
||||
async for line in response.content:
|
||||
line = line.decode('utf-8').strip()
|
||||
|
||||
if not line:
|
||||
continue
|
||||
|
||||
if line.startswith('data: '):
|
||||
data = line[6:] # Remove 'data: ' prefix
|
||||
|
||||
if data == '[DONE]':
|
||||
break
|
||||
|
||||
try:
|
||||
import json
|
||||
chunk_data = json.loads(data)
|
||||
|
||||
# Extract text from chunk
|
||||
if 'choices' in chunk_data and len(chunk_data['choices']) > 0:
|
||||
choice = chunk_data['choices'][0]
|
||||
if 'delta' in choice and 'content' in choice['delta']:
|
||||
content = choice['delta']['content']
|
||||
if content:
|
||||
yield LlmChunk(
|
||||
text=content,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=False
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Failed to parse chunk: {data}")
|
||||
continue
|
||||
|
||||
# Send final chunk
|
||||
yield LlmChunk(
|
||||
text="",
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=True
|
||||
)
|
||||
|
||||
logger.debug("Streaming complete")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"TGI streaming exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import logging
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import LlmService, LlmResult
|
||||
from .... base import LlmService, LlmResult, LlmChunk
|
||||
|
||||
default_ident = "text-completion"
|
||||
|
||||
|
|
@ -113,6 +113,89 @@ class Processor(LlmService):
|
|||
logger.error(f"vLLM LLM exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
def supports_streaming(self):
|
||||
"""vLLM supports streaming"""
|
||||
return True
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
||||
"""Stream content generation from vLLM"""
|
||||
model_name = model or self.default_model
|
||||
effective_temperature = temperature if temperature is not None else self.temperature
|
||||
|
||||
logger.debug(f"Using model (streaming): {model_name}")
|
||||
logger.debug(f"Using temperature: {effective_temperature}")
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
request = {
|
||||
"model": model_name,
|
||||
"prompt": system + "\n\n" + prompt,
|
||||
"max_tokens": self.max_output,
|
||||
"temperature": effective_temperature,
|
||||
"stream": True,
|
||||
}
|
||||
|
||||
try:
|
||||
url = f"{self.base_url}/completions"
|
||||
|
||||
async with self.session.post(
|
||||
url,
|
||||
headers=headers,
|
||||
json=request,
|
||||
) as response:
|
||||
|
||||
if response.status != 200:
|
||||
raise RuntimeError("Bad status: " + str(response.status))
|
||||
|
||||
# Parse SSE stream
|
||||
async for line in response.content:
|
||||
line = line.decode('utf-8').strip()
|
||||
|
||||
if not line:
|
||||
continue
|
||||
|
||||
if line.startswith('data: '):
|
||||
data = line[6:] # Remove 'data: ' prefix
|
||||
|
||||
if data == '[DONE]':
|
||||
break
|
||||
|
||||
try:
|
||||
import json
|
||||
chunk_data = json.loads(data)
|
||||
|
||||
# Extract text from chunk
|
||||
if 'choices' in chunk_data and len(chunk_data['choices']) > 0:
|
||||
choice = chunk_data['choices'][0]
|
||||
if 'text' in choice and choice['text']:
|
||||
yield LlmChunk(
|
||||
text=choice['text'],
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=False
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Failed to parse chunk: {data}")
|
||||
continue
|
||||
|
||||
# Send final chunk
|
||||
yield LlmChunk(
|
||||
text="",
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=True
|
||||
)
|
||||
|
||||
logger.debug("Streaming complete")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"vLLM streaming exception ({type(e).__name__}): {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -101,6 +101,9 @@ class Processor(FlowProcessor):
|
|||
|
||||
kind = v.id
|
||||
|
||||
# Check if streaming is requested
|
||||
streaming = getattr(v, 'streaming', False)
|
||||
|
||||
try:
|
||||
|
||||
logger.debug(f"Prompt terms: {v.terms}")
|
||||
|
|
@ -109,16 +112,68 @@ class Processor(FlowProcessor):
|
|||
k: json.loads(v)
|
||||
for k, v in v.terms.items()
|
||||
}
|
||||
|
||||
logger.debug(f"Handling prompt kind {kind}...")
|
||||
|
||||
logger.debug(f"Handling prompt kind {kind}... (streaming={streaming})")
|
||||
|
||||
# If streaming, we need to handle it differently
|
||||
if streaming:
|
||||
# For streaming, we need to intercept LLM responses
|
||||
# and forward them as they arrive
|
||||
|
||||
async def llm_streaming(system, prompt):
|
||||
logger.debug(f"System prompt: {system}")
|
||||
logger.debug(f"User prompt: {prompt}")
|
||||
|
||||
# Use the text completion client with recipient handler
|
||||
client = flow("text-completion-request")
|
||||
|
||||
async def forward_chunks(resp):
|
||||
if resp.error:
|
||||
raise RuntimeError(resp.error.message)
|
||||
|
||||
is_final = getattr(resp, 'end_of_stream', False)
|
||||
|
||||
# Always send a message if there's content OR if it's the final message
|
||||
if resp.response or is_final:
|
||||
# Forward each chunk immediately
|
||||
r = PromptResponse(
|
||||
text=resp.response if resp.response else "",
|
||||
object=None,
|
||||
error=None,
|
||||
end_of_stream=is_final,
|
||||
)
|
||||
await flow("response").send(r, properties={"id": id})
|
||||
|
||||
# Return True when end_of_stream
|
||||
return is_final
|
||||
|
||||
await client.request(
|
||||
TextCompletionRequest(
|
||||
system=system, prompt=prompt, streaming=True
|
||||
),
|
||||
recipient=forward_chunks,
|
||||
timeout=600
|
||||
)
|
||||
|
||||
# Return empty string since we already sent all chunks
|
||||
return ""
|
||||
|
||||
try:
|
||||
await self.manager.invoke(kind, input, llm_streaming)
|
||||
except Exception as e:
|
||||
logger.error(f"Prompt streaming exception: {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
return
|
||||
|
||||
# Non-streaming path (original behavior)
|
||||
async def llm(system, prompt):
|
||||
|
||||
logger.debug(f"System prompt: {system}")
|
||||
logger.debug(f"User prompt: {prompt}")
|
||||
|
||||
resp = await flow("text-completion-request").text_completion(
|
||||
system = system, prompt = prompt,
|
||||
system = system, prompt = prompt, streaming = False,
|
||||
)
|
||||
|
||||
try:
|
||||
|
|
@ -143,6 +198,7 @@ class Processor(FlowProcessor):
|
|||
text=resp,
|
||||
object=None,
|
||||
error=None,
|
||||
end_of_stream=True,
|
||||
)
|
||||
|
||||
await flow("response").send(r, properties={"id": id})
|
||||
|
|
@ -158,6 +214,7 @@ class Processor(FlowProcessor):
|
|||
text=None,
|
||||
object=json.dumps(resp),
|
||||
error=None,
|
||||
end_of_stream=True,
|
||||
)
|
||||
|
||||
await flow("response").send(r, properties={"id": id})
|
||||
|
|
@ -175,27 +232,13 @@ class Processor(FlowProcessor):
|
|||
type = "llm-error",
|
||||
message = str(e),
|
||||
),
|
||||
response=None,
|
||||
text=None,
|
||||
object=None,
|
||||
end_of_stream=True,
|
||||
)
|
||||
|
||||
await flow("response").send(r, properties={"id": id})
|
||||
|
||||
except Exception as e:
|
||||
|
||||
logger.error(f"Prompt service exception: {e}", exc_info=True)
|
||||
|
||||
logger.debug("Sending error response...")
|
||||
|
||||
r = PromptResponse(
|
||||
error=Error(
|
||||
type = "llm-error",
|
||||
message = str(e),
|
||||
),
|
||||
response=None,
|
||||
)
|
||||
|
||||
await self.send(r, properties={"id": id})
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -32,7 +32,7 @@ from vertexai.generative_models import (
|
|||
from anthropic import AnthropicVertex, RateLimitError
|
||||
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import LlmService, LlmResult
|
||||
from .... base import LlmService, LlmResult, LlmChunk
|
||||
|
||||
# Module logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -239,6 +239,123 @@ class Processor(LlmService):
|
|||
logger.error(f"VertexAI LLM exception: {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
def supports_streaming(self):
|
||||
"""VertexAI supports streaming for both Gemini and Claude models"""
|
||||
return True
|
||||
|
||||
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
||||
"""
|
||||
Stream content generation from VertexAI (Gemini or Claude).
|
||||
Yields LlmChunk objects with is_final=True on the last chunk.
|
||||
"""
|
||||
# Use provided model or fall back to default
|
||||
model_name = model or self.default_model
|
||||
# Use provided temperature or fall back to default
|
||||
effective_temperature = temperature if temperature is not None else self.temperature
|
||||
|
||||
logger.debug(f"Using model (streaming): {model_name}")
|
||||
logger.debug(f"Using temperature: {effective_temperature}")
|
||||
|
||||
try:
|
||||
if 'claude' in model_name.lower():
|
||||
# Claude/Anthropic streaming
|
||||
logger.debug(f"Streaming request to Anthropic model '{model_name}'...")
|
||||
client = self._get_anthropic_client()
|
||||
|
||||
total_in_tokens = 0
|
||||
total_out_tokens = 0
|
||||
|
||||
with client.messages.stream(
|
||||
model=model_name,
|
||||
system=system,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
max_tokens=self.api_params['max_output_tokens'],
|
||||
temperature=effective_temperature,
|
||||
top_p=self.api_params['top_p'],
|
||||
top_k=self.api_params['top_k'],
|
||||
) as stream:
|
||||
# Stream text chunks
|
||||
for text in stream.text_stream:
|
||||
yield LlmChunk(
|
||||
text=text,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=False
|
||||
)
|
||||
|
||||
# Get final message with token counts
|
||||
final_message = stream.get_final_message()
|
||||
total_in_tokens = final_message.usage.input_tokens
|
||||
total_out_tokens = final_message.usage.output_tokens
|
||||
|
||||
# Send final chunk with token counts
|
||||
yield LlmChunk(
|
||||
text="",
|
||||
in_token=total_in_tokens,
|
||||
out_token=total_out_tokens,
|
||||
model=model_name,
|
||||
is_final=True
|
||||
)
|
||||
|
||||
logger.info(f"Input Tokens: {total_in_tokens}")
|
||||
logger.info(f"Output Tokens: {total_out_tokens}")
|
||||
|
||||
else:
|
||||
# Gemini streaming
|
||||
logger.debug(f"Streaming request to Gemini model '{model_name}'...")
|
||||
full_prompt = system + "\n\n" + prompt
|
||||
|
||||
llm, generation_config = self._get_gemini_model(model_name, effective_temperature)
|
||||
|
||||
response = llm.generate_content(
|
||||
full_prompt,
|
||||
generation_config=generation_config,
|
||||
safety_settings=self.safety_settings,
|
||||
stream=True # Enable streaming
|
||||
)
|
||||
|
||||
total_in_tokens = 0
|
||||
total_out_tokens = 0
|
||||
|
||||
# Stream chunks
|
||||
for chunk in response:
|
||||
if chunk.text:
|
||||
yield LlmChunk(
|
||||
text=chunk.text,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=model_name,
|
||||
is_final=False
|
||||
)
|
||||
|
||||
# Accumulate token counts if available
|
||||
if hasattr(chunk, 'usage_metadata') and chunk.usage_metadata:
|
||||
if hasattr(chunk.usage_metadata, 'prompt_token_count'):
|
||||
total_in_tokens = chunk.usage_metadata.prompt_token_count
|
||||
if hasattr(chunk.usage_metadata, 'candidates_token_count'):
|
||||
total_out_tokens = chunk.usage_metadata.candidates_token_count
|
||||
|
||||
# Send final chunk with token counts
|
||||
yield LlmChunk(
|
||||
text="",
|
||||
in_token=total_in_tokens,
|
||||
out_token=total_out_tokens,
|
||||
model=model_name,
|
||||
is_final=True
|
||||
)
|
||||
|
||||
logger.info(f"Input Tokens: {total_in_tokens}")
|
||||
logger.info(f"Output Tokens: {total_out_tokens}")
|
||||
|
||||
except (google.api_core.exceptions.ResourceExhausted, RateLimitError) as e:
|
||||
logger.warning(f"Hit rate limit during streaming: {e}")
|
||||
raise TooManyRequests()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"VertexAI streaming exception: {e}", exc_info=True)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue