diff --git a/trustgraph-base/trustgraph/api/async_metrics.py b/trustgraph-base/trustgraph/api/async_metrics.py index 3a00926f..9ba22f02 100644 --- a/trustgraph-base/trustgraph/api/async_metrics.py +++ b/trustgraph-base/trustgraph/api/async_metrics.py @@ -1,21 +1,21 @@ import aiohttp -from typing import Optional +from typing import Optional, Dict class AsyncMetrics: """Asynchronous metrics client""" - def __init__(self, url: str, timeout: int, token: Optional[str]): - self.url = url - self.timeout = timeout - self.token = token + def __init__(self, url: str, timeout: int, token: Optional[str]) -> None: + self.url: str = url + self.timeout: int = timeout + self.token: Optional[str] = token async def get(self) -> str: """Get Prometheus metrics as text""" - url = f"{self.url}/api/metrics" + url: str = f"{self.url}/api/metrics" - headers = {} + headers: Dict[str, str] = {} if self.token: headers["Authorization"] = f"Bearer {self.token}" @@ -27,3 +27,7 @@ class AsyncMetrics: raise Exception(f"Status code {resp.status}") return await resp.text() + + async def aclose(self) -> None: + """Close connections""" + pass diff --git a/trustgraph-base/trustgraph/api/metrics.py b/trustgraph-base/trustgraph/api/metrics.py index 4dcd93e0..68968349 100644 --- a/trustgraph-base/trustgraph/api/metrics.py +++ b/trustgraph-base/trustgraph/api/metrics.py @@ -1,21 +1,21 @@ import requests -from typing import Optional +from typing import Optional, Dict class Metrics: """Synchronous metrics client""" - def __init__(self, url: str, timeout: int, token: Optional[str]): - self.url = url - self.timeout = timeout - self.token = token + def __init__(self, url: str, timeout: int, token: Optional[str]) -> None: + self.url: str = url + self.timeout: int = timeout + self.token: Optional[str] = token def get(self) -> str: """Get Prometheus metrics as text""" - url = f"{self.url}/api/metrics" + url: str = f"{self.url}/api/metrics" - headers = {} + headers: Dict[str, str] = {} if self.token: headers["Authorization"] = f"Bearer {self.token}" diff --git a/trustgraph-base/trustgraph/api/socket_client.py b/trustgraph-base/trustgraph/api/socket_client.py index 3e318fb7..7e5dfaaf 100644 --- a/trustgraph-base/trustgraph/api/socket_client.py +++ b/trustgraph-base/trustgraph/api/socket_client.py @@ -2,24 +2,24 @@ import json import asyncio import websockets -from typing import Optional, Dict, Any, Iterator, Union +from typing import Optional, Dict, Any, Iterator, Union, List from threading import Lock -from . types import AgentThought, AgentObservation, AgentAnswer, RAGChunk +from . types import AgentThought, AgentObservation, AgentAnswer, RAGChunk, StreamingChunk from . exceptions import ProtocolException, ApplicationException class SocketClient: """Synchronous WebSocket client (wraps async websockets library)""" - def __init__(self, url: str, timeout: int, token: Optional[str]): - self.url = self._convert_to_ws_url(url) - self.timeout = timeout - self.token = token - self._connection = None - self._request_counter = 0 - self._lock = Lock() - self._loop = None + def __init__(self, url: str, timeout: int, token: Optional[str]) -> None: + self.url: str = self._convert_to_ws_url(url) + self.timeout: int = timeout + self.token: Optional[str] = token + self._connection: Optional[Any] = None + self._request_counter: int = 0 + self._lock: Lock = Lock() + self._loop: Optional[asyncio.AbstractEventLoop] = None def _convert_to_ws_url(self, url: str) -> str: """Convert HTTP URL to WebSocket URL""" @@ -33,11 +33,17 @@ class SocketClient: # Assume ws:// return f"ws://{url}" - def flow(self, flow_id: str): + def flow(self, flow_id: str) -> "SocketFlowInstance": """Get flow instance for WebSocket operations""" return SocketFlowInstance(self, flow_id) - def _send_request_sync(self, service: str, flow: Optional[str], request: Dict[str, Any], streaming: bool = False): + def _send_request_sync( + self, + service: str, + flow: Optional[str], + request: Dict[str, Any], + streaming: bool = False + ) -> Union[Dict[str, Any], Iterator[StreamingChunk]]: """Synchronous wrapper around async WebSocket communication""" # Create event loop if needed try: @@ -58,7 +64,13 @@ class SocketClient: # For non-streaming, just run the async code and return result return loop.run_until_complete(self._send_request_async(service, flow, request)) - def _streaming_generator(self, service: str, flow: Optional[str], request: Dict[str, Any], loop): + def _streaming_generator( + self, + service: str, + flow: Optional[str], + request: Dict[str, Any], + loop: asyncio.AbstractEventLoop + ) -> Iterator[StreamingChunk]: """Generator that yields streaming chunks""" async_gen = self._send_request_async_streaming(service, flow, request) @@ -76,7 +88,12 @@ class SocketClient: except: pass - async def _send_request_async(self, service: str, flow: Optional[str], request: Dict[str, Any]): + async def _send_request_async( + self, + service: str, + flow: Optional[str], + request: Dict[str, Any] + ) -> Dict[str, Any]: """Async implementation of WebSocket request (non-streaming)""" # Generate unique request ID with self._lock: @@ -116,7 +133,12 @@ class SocketClient: return response["response"] - async def _send_request_async_streaming(self, service: str, flow: Optional[str], request: Dict[str, Any]): + async def _send_request_async_streaming( + self, + service: str, + flow: Optional[str], + request: Dict[str, Any] + ) -> Iterator[StreamingChunk]: """Async implementation of WebSocket request (streaming)""" # Generate unique request ID with self._lock: @@ -162,7 +184,7 @@ class SocketClient: if resp.get("end_of_stream") or resp.get("end_of_dialog") or response.get("complete"): break - def _parse_chunk(self, resp: Dict[str, Any]): + def _parse_chunk(self, resp: Dict[str, Any]) -> StreamingChunk: """Parse response chunk into appropriate type""" chunk_type = resp.get("chunk_type") @@ -190,7 +212,7 @@ class SocketClient: error=resp.get("error") ) - def close(self): + def close(self) -> None: """Close WebSocket connection""" # Cleanup handled by context manager in async code pass @@ -199,13 +221,20 @@ class SocketClient: class SocketFlowInstance: """Synchronous WebSocket flow instance with same interface as REST FlowInstance""" - def __init__(self, client: SocketClient, flow_id: str): - self.client = client - self.flow_id = flow_id + def __init__(self, client: SocketClient, flow_id: str) -> None: + self.client: SocketClient = client + self.flow_id: str = flow_id - def agent(self, question: str, user: str, state: Optional[Dict[str, Any]] = None, - group: Optional[str] = None, history: Optional[list] = None, - streaming: bool = False, **kwargs) -> Union[Dict[str, Any], Iterator]: + def agent( + self, + question: str, + user: str, + state: Optional[Dict[str, Any]] = None, + group: Optional[str] = None, + history: Optional[List[Dict[str, Any]]] = None, + streaming: bool = False, + **kwargs: Any + ) -> Union[Dict[str, Any], Iterator[StreamingChunk]]: """Agent with optional streaming""" request = { "question": question, @@ -241,9 +270,17 @@ class SocketFlowInstance: else: return result.get("response", "") - def graph_rag(self, question: str, user: str, collection: str, - max_subgraph_size: int = 1000, max_subgraph_count: int = 5, - max_entity_distance: int = 3, streaming: bool = False, **kwargs) -> Union[str, Iterator[str]]: + def graph_rag( + self, + question: str, + user: str, + collection: str, + max_subgraph_size: int = 1000, + max_subgraph_count: int = 5, + max_entity_distance: int = 3, + streaming: bool = False, + **kwargs: Any + ) -> Union[str, Iterator[str]]: """Graph RAG with optional streaming""" request = { "question": question, @@ -265,8 +302,15 @@ class SocketFlowInstance: else: return result.get("response", "") - def document_rag(self, question: str, user: str, collection: str, - doc_limit: int = 10, streaming: bool = False, **kwargs) -> Union[str, Iterator[str]]: + def document_rag( + self, + question: str, + user: str, + collection: str, + doc_limit: int = 10, + streaming: bool = False, + **kwargs: Any + ) -> Union[str, Iterator[str]]: """Document RAG with optional streaming""" request = { "question": question, @@ -286,7 +330,13 @@ class SocketFlowInstance: else: return result.get("response", "") - def prompt(self, id: str, variables: Dict[str, str], streaming: bool = False, **kwargs) -> Union[str, Iterator[str]]: + def prompt( + self, + id: str, + variables: Dict[str, str], + streaming: bool = False, + **kwargs: Any + ) -> Union[str, Iterator[str]]: """Execute prompt with optional streaming""" request = { "id": id, @@ -304,7 +354,14 @@ class SocketFlowInstance: else: return result.get("response", "") - def graph_embeddings_query(self, text: str, user: str, collection: str, limit: int = 10, **kwargs): + def graph_embeddings_query( + self, + text: str, + user: str, + collection: str, + limit: int = 10, + **kwargs: Any + ) -> Dict[str, Any]: """Query graph embeddings for semantic search""" request = { "text": text, @@ -316,14 +373,23 @@ class SocketFlowInstance: return self.client._send_request_sync("graph-embeddings", self.flow_id, request, False) - def embeddings(self, text: str, **kwargs): + def embeddings(self, text: str, **kwargs: Any) -> Dict[str, Any]: """Generate text embeddings""" request = {"text": text} request.update(kwargs) return self.client._send_request_sync("embeddings", self.flow_id, request, False) - def triples_query(self, s=None, p=None, o=None, user=None, collection=None, limit=100, **kwargs): + def triples_query( + self, + s: Optional[str] = None, + p: Optional[str] = None, + o: Optional[str] = None, + user: Optional[str] = None, + collection: Optional[str] = None, + limit: int = 100, + **kwargs: Any + ) -> Dict[str, Any]: """Triple pattern query""" request = {"limit": limit} if s is not None: @@ -340,8 +406,15 @@ class SocketFlowInstance: return self.client._send_request_sync("triples", self.flow_id, request, False) - def objects_query(self, query: str, user: str, collection: str, variables: Optional[Dict] = None, - operation_name: Optional[str] = None, **kwargs): + def objects_query( + self, + query: str, + user: str, + collection: str, + variables: Optional[Dict[str, Any]] = None, + operation_name: Optional[str] = None, + **kwargs: Any + ) -> Dict[str, Any]: """GraphQL query""" request = { "query": query, @@ -356,7 +429,12 @@ class SocketFlowInstance: return self.client._send_request_sync("objects", self.flow_id, request, False) - def mcp_tool(self, name: str, parameters: Dict[str, Any], **kwargs): + def mcp_tool( + self, + name: str, + parameters: Dict[str, Any], + **kwargs: Any + ) -> Dict[str, Any]: """Execute MCP tool""" request = { "name": name,