mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-24 20:51:02 +02:00
Type hints
This commit is contained in:
Cyber MacGeddon
parent
9e5afef1b3
commit
3cdb4cd74b
3 changed files with 53 additions and 53 deletions
|
|
@ -1,7 +1,7 @@
|
|||
|
||||
import json
|
||||
import websockets
|
||||
from typing import Optional, AsyncIterator, Dict, Any
|
||||
from typing import Optional, AsyncIterator, Dict, Any, Iterator
|
||||
|
||||
from . types import Triple
|
||||
|
||||
|
|
@ -9,10 +9,10 @@ from . types import Triple
|
|||
class AsyncBulkClient:
|
||||
"""Asynchronous bulk operations client"""
|
||||
|
||||
def __init__(self, url: str, timeout: int, token: Optional[str]):
|
||||
self.url = self._convert_to_ws_url(url)
|
||||
self.timeout = timeout
|
||||
self.token = token
|
||||
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
|
||||
|
||||
def _convert_to_ws_url(self, url: str) -> str:
|
||||
"""Convert HTTP URL to WebSocket URL"""
|
||||
|
|
@ -25,7 +25,7 @@ class AsyncBulkClient:
|
|||
else:
|
||||
return f"ws://{url}"
|
||||
|
||||
async def import_triples(self, flow: str, triples: AsyncIterator[Triple], **kwargs) -> None:
|
||||
async def import_triples(self, flow: str, triples: AsyncIterator[Triple], **kwargs: Any) -> None:
|
||||
"""Bulk import triples via WebSocket"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/import/triples"
|
||||
if self.token:
|
||||
|
|
@ -40,7 +40,7 @@ class AsyncBulkClient:
|
|||
}
|
||||
await websocket.send(json.dumps(message))
|
||||
|
||||
async def export_triples(self, flow: str, **kwargs) -> AsyncIterator[Triple]:
|
||||
async def export_triples(self, flow: str, **kwargs: Any) -> AsyncIterator[Triple]:
|
||||
"""Bulk export triples via WebSocket"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/export/triples"
|
||||
if self.token:
|
||||
|
|
@ -55,7 +55,7 @@ class AsyncBulkClient:
|
|||
o=data.get("o", "")
|
||||
)
|
||||
|
||||
async def import_graph_embeddings(self, flow: str, embeddings: AsyncIterator[Dict[str, Any]], **kwargs) -> None:
|
||||
async def import_graph_embeddings(self, flow: str, embeddings: AsyncIterator[Dict[str, Any]], **kwargs: Any) -> None:
|
||||
"""Bulk import graph embeddings via WebSocket"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/import/graph-embeddings"
|
||||
if self.token:
|
||||
|
|
@ -65,7 +65,7 @@ class AsyncBulkClient:
|
|||
async for embedding in embeddings:
|
||||
await websocket.send(json.dumps(embedding))
|
||||
|
||||
async def export_graph_embeddings(self, flow: str, **kwargs) -> AsyncIterator[Dict[str, Any]]:
|
||||
async def export_graph_embeddings(self, flow: str, **kwargs: Any) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""Bulk export graph embeddings via WebSocket"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/export/graph-embeddings"
|
||||
if self.token:
|
||||
|
|
@ -75,7 +75,7 @@ class AsyncBulkClient:
|
|||
async for raw_message in websocket:
|
||||
yield json.loads(raw_message)
|
||||
|
||||
async def import_document_embeddings(self, flow: str, embeddings: AsyncIterator[Dict[str, Any]], **kwargs) -> None:
|
||||
async def import_document_embeddings(self, flow: str, embeddings: AsyncIterator[Dict[str, Any]], **kwargs: Any) -> None:
|
||||
"""Bulk import document embeddings via WebSocket"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/import/document-embeddings"
|
||||
if self.token:
|
||||
|
|
@ -85,7 +85,7 @@ class AsyncBulkClient:
|
|||
async for embedding in embeddings:
|
||||
await websocket.send(json.dumps(embedding))
|
||||
|
||||
async def export_document_embeddings(self, flow: str, **kwargs) -> AsyncIterator[Dict[str, Any]]:
|
||||
async def export_document_embeddings(self, flow: str, **kwargs: Any) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""Bulk export document embeddings via WebSocket"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/export/document-embeddings"
|
||||
if self.token:
|
||||
|
|
@ -95,7 +95,7 @@ class AsyncBulkClient:
|
|||
async for raw_message in websocket:
|
||||
yield json.loads(raw_message)
|
||||
|
||||
async def import_entity_contexts(self, flow: str, contexts: AsyncIterator[Dict[str, Any]], **kwargs) -> None:
|
||||
async def import_entity_contexts(self, flow: str, contexts: AsyncIterator[Dict[str, Any]], **kwargs: Any) -> None:
|
||||
"""Bulk import entity contexts via WebSocket"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/import/entity-contexts"
|
||||
if self.token:
|
||||
|
|
@ -105,7 +105,7 @@ class AsyncBulkClient:
|
|||
async for context in contexts:
|
||||
await websocket.send(json.dumps(context))
|
||||
|
||||
async def export_entity_contexts(self, flow: str, **kwargs) -> AsyncIterator[Dict[str, Any]]:
|
||||
async def export_entity_contexts(self, flow: str, **kwargs: Any) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""Bulk export entity contexts via WebSocket"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/export/entity-contexts"
|
||||
if self.token:
|
||||
|
|
@ -115,7 +115,7 @@ class AsyncBulkClient:
|
|||
async for raw_message in websocket:
|
||||
yield json.loads(raw_message)
|
||||
|
||||
async def import_objects(self, flow: str, objects: AsyncIterator[Dict[str, Any]], **kwargs) -> None:
|
||||
async def import_objects(self, flow: str, objects: AsyncIterator[Dict[str, Any]], **kwargs: Any) -> None:
|
||||
"""Bulk import objects via WebSocket"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/import/objects"
|
||||
if self.token:
|
||||
|
|
@ -125,7 +125,7 @@ class AsyncBulkClient:
|
|||
async for obj in objects:
|
||||
await websocket.send(json.dumps(obj))
|
||||
|
||||
async def aclose(self):
|
||||
async def aclose(self) -> None:
|
||||
"""Close connections"""
|
||||
# Cleanup handled by context managers
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -20,10 +20,10 @@ def check_error(response):
|
|||
class AsyncFlow:
|
||||
"""Asynchronous REST-based flow interface"""
|
||||
|
||||
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 request(self, path: str, request_data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Make async HTTP request to Gateway API"""
|
||||
|
|
@ -113,7 +113,7 @@ class AsyncFlow:
|
|||
"""Get async flow instance"""
|
||||
return AsyncFlowInstance(self, flow_id)
|
||||
|
||||
async def aclose(self):
|
||||
async def aclose(self) -> None:
|
||||
"""Close connection (cleanup handled by aiohttp session)"""
|
||||
pass
|
||||
|
||||
|
|
@ -130,7 +130,7 @@ class AsyncFlowInstance:
|
|||
return await self.flow.request(f"flow/{self.flow_id}/service/{service}", request_data)
|
||||
|
||||
async def agent(self, question: str, user: str, state: Optional[Dict] = None,
|
||||
group: Optional[str] = None, history: Optional[List] = None, **kwargs) -> Dict[str, Any]:
|
||||
group: Optional[str] = None, history: Optional[List] = None, **kwargs: Any) -> Dict[str, Any]:
|
||||
"""Execute agent (non-streaming, use async_socket for streaming)"""
|
||||
request_data = {
|
||||
"question": question,
|
||||
|
|
@ -147,7 +147,7 @@ class AsyncFlowInstance:
|
|||
|
||||
return await self.request("agent", request_data)
|
||||
|
||||
async def text_completion(self, system: str, prompt: str, **kwargs) -> str:
|
||||
async def text_completion(self, system: str, prompt: str, **kwargs: Any) -> str:
|
||||
"""Text completion (non-streaming, use async_socket for streaming)"""
|
||||
request_data = {
|
||||
"system": system,
|
||||
|
|
@ -161,7 +161,7 @@ class AsyncFlowInstance:
|
|||
|
||||
async 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, **kwargs) -> str:
|
||||
max_entity_distance: int = 3, **kwargs: Any) -> str:
|
||||
"""Graph RAG (non-streaming, use async_socket for streaming)"""
|
||||
request_data = {
|
||||
"question": question,
|
||||
|
|
@ -178,7 +178,7 @@ class AsyncFlowInstance:
|
|||
return result.get("response", "")
|
||||
|
||||
async def document_rag(self, question: str, user: str, collection: str,
|
||||
doc_limit: int = 10, **kwargs) -> str:
|
||||
doc_limit: int = 10, **kwargs: Any) -> str:
|
||||
"""Document RAG (non-streaming, use async_socket for streaming)"""
|
||||
request_data = {
|
||||
"question": question,
|
||||
|
|
@ -192,7 +192,7 @@ class AsyncFlowInstance:
|
|||
result = await self.request("document-rag", request_data)
|
||||
return result.get("response", "")
|
||||
|
||||
async def graph_embeddings_query(self, text: str, user: str, collection: str, limit: int = 10, **kwargs):
|
||||
async def graph_embeddings_query(self, text: str, user: str, collection: str, limit: int = 10, **kwargs: Any):
|
||||
"""Query graph embeddings for semantic search"""
|
||||
request_data = {
|
||||
"text": text,
|
||||
|
|
@ -204,14 +204,14 @@ class AsyncFlowInstance:
|
|||
|
||||
return await self.request("graph-embeddings", request_data)
|
||||
|
||||
async def embeddings(self, text: str, **kwargs):
|
||||
async def embeddings(self, text: str, **kwargs: Any):
|
||||
"""Generate text embeddings"""
|
||||
request_data = {"text": text}
|
||||
request_data.update(kwargs)
|
||||
|
||||
return await self.request("embeddings", request_data)
|
||||
|
||||
async def triples_query(self, s=None, p=None, o=None, user=None, collection=None, limit=100, **kwargs):
|
||||
async def triples_query(self, s=None, p=None, o=None, user=None, collection=None, limit=100, **kwargs: Any):
|
||||
"""Triple pattern query"""
|
||||
request_data = {"limit": limit}
|
||||
if s is not None:
|
||||
|
|
@ -229,7 +229,7 @@ class AsyncFlowInstance:
|
|||
return await self.request("triples", request_data)
|
||||
|
||||
async def objects_query(self, query: str, user: str, collection: str, variables: Optional[Dict] = None,
|
||||
operation_name: Optional[str] = None, **kwargs):
|
||||
operation_name: Optional[str] = None, **kwargs: Any):
|
||||
"""GraphQL query"""
|
||||
request_data = {
|
||||
"query": query,
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
import json
|
||||
import asyncio
|
||||
import websockets
|
||||
from typing import Optional, Iterator, Dict, Any
|
||||
from typing import Optional, Iterator, Dict, Any, Coroutine
|
||||
|
||||
from . types import Triple
|
||||
from . exceptions import ProtocolException
|
||||
|
|
@ -11,10 +11,10 @@ from . exceptions import ProtocolException
|
|||
class BulkClient:
|
||||
"""Synchronous bulk operations client"""
|
||||
|
||||
def __init__(self, url: str, timeout: int, token: Optional[str]):
|
||||
self.url = self._convert_to_ws_url(url)
|
||||
self.timeout = timeout
|
||||
self.token = token
|
||||
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
|
||||
|
||||
def _convert_to_ws_url(self, url: str) -> str:
|
||||
"""Convert HTTP URL to WebSocket URL"""
|
||||
|
|
@ -27,7 +27,7 @@ class BulkClient:
|
|||
else:
|
||||
return f"ws://{url}"
|
||||
|
||||
def _run_async(self, coro):
|
||||
def _run_async(self, coro: Coroutine[Any, Any, Any]) -> Any:
|
||||
"""Run async coroutine synchronously"""
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
|
|
@ -40,11 +40,11 @@ class BulkClient:
|
|||
|
||||
return loop.run_until_complete(coro)
|
||||
|
||||
def import_triples(self, flow: str, triples: Iterator[Triple], **kwargs) -> None:
|
||||
def import_triples(self, flow: str, triples: Iterator[Triple], **kwargs: Any) -> None:
|
||||
"""Bulk import triples via WebSocket"""
|
||||
self._run_async(self._import_triples_async(flow, triples))
|
||||
|
||||
async def _import_triples_async(self, flow: str, triples: Iterator[Triple]):
|
||||
async def _import_triples_async(self, flow: str, triples: Iterator[Triple]) -> None:
|
||||
"""Async implementation of triple import"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/import/triples"
|
||||
if self.token:
|
||||
|
|
@ -59,7 +59,7 @@ class BulkClient:
|
|||
}
|
||||
await websocket.send(json.dumps(message))
|
||||
|
||||
def export_triples(self, flow: str, **kwargs) -> Iterator[Triple]:
|
||||
def export_triples(self, flow: str, **kwargs: Any) -> Iterator[Triple]:
|
||||
"""Bulk export triples via WebSocket"""
|
||||
async_gen = self._export_triples_async(flow)
|
||||
|
||||
|
|
@ -85,7 +85,7 @@ class BulkClient:
|
|||
except:
|
||||
pass
|
||||
|
||||
async def _export_triples_async(self, flow: str):
|
||||
async def _export_triples_async(self, flow: str) -> Iterator[Triple]:
|
||||
"""Async implementation of triple export"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/export/triples"
|
||||
if self.token:
|
||||
|
|
@ -100,11 +100,11 @@ class BulkClient:
|
|||
o=data.get("o", "")
|
||||
)
|
||||
|
||||
def import_graph_embeddings(self, flow: str, embeddings: Iterator[Dict[str, Any]], **kwargs) -> None:
|
||||
def import_graph_embeddings(self, flow: str, embeddings: Iterator[Dict[str, Any]], **kwargs: Any) -> None:
|
||||
"""Bulk import graph embeddings via WebSocket"""
|
||||
self._run_async(self._import_graph_embeddings_async(flow, embeddings))
|
||||
|
||||
async def _import_graph_embeddings_async(self, flow: str, embeddings: Iterator[Dict[str, Any]]):
|
||||
async def _import_graph_embeddings_async(self, flow: str, embeddings: Iterator[Dict[str, Any]]) -> None:
|
||||
"""Async implementation of graph embeddings import"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/import/graph-embeddings"
|
||||
if self.token:
|
||||
|
|
@ -114,7 +114,7 @@ class BulkClient:
|
|||
for embedding in embeddings:
|
||||
await websocket.send(json.dumps(embedding))
|
||||
|
||||
def export_graph_embeddings(self, flow: str, **kwargs) -> Iterator[Dict[str, Any]]:
|
||||
def export_graph_embeddings(self, flow: str, **kwargs: Any) -> Iterator[Dict[str, Any]]:
|
||||
"""Bulk export graph embeddings via WebSocket"""
|
||||
async_gen = self._export_graph_embeddings_async(flow)
|
||||
|
||||
|
|
@ -140,7 +140,7 @@ class BulkClient:
|
|||
except:
|
||||
pass
|
||||
|
||||
async def _export_graph_embeddings_async(self, flow: str):
|
||||
async def _export_graph_embeddings_async(self, flow: str) -> Iterator[Dict[str, Any]]:
|
||||
"""Async implementation of graph embeddings export"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/export/graph-embeddings"
|
||||
if self.token:
|
||||
|
|
@ -150,11 +150,11 @@ class BulkClient:
|
|||
async for raw_message in websocket:
|
||||
yield json.loads(raw_message)
|
||||
|
||||
def import_document_embeddings(self, flow: str, embeddings: Iterator[Dict[str, Any]], **kwargs) -> None:
|
||||
def import_document_embeddings(self, flow: str, embeddings: Iterator[Dict[str, Any]], **kwargs: Any) -> None:
|
||||
"""Bulk import document embeddings via WebSocket"""
|
||||
self._run_async(self._import_document_embeddings_async(flow, embeddings))
|
||||
|
||||
async def _import_document_embeddings_async(self, flow: str, embeddings: Iterator[Dict[str, Any]]):
|
||||
async def _import_document_embeddings_async(self, flow: str, embeddings: Iterator[Dict[str, Any]]) -> None:
|
||||
"""Async implementation of document embeddings import"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/import/document-embeddings"
|
||||
if self.token:
|
||||
|
|
@ -164,7 +164,7 @@ class BulkClient:
|
|||
for embedding in embeddings:
|
||||
await websocket.send(json.dumps(embedding))
|
||||
|
||||
def export_document_embeddings(self, flow: str, **kwargs) -> Iterator[Dict[str, Any]]:
|
||||
def export_document_embeddings(self, flow: str, **kwargs: Any) -> Iterator[Dict[str, Any]]:
|
||||
"""Bulk export document embeddings via WebSocket"""
|
||||
async_gen = self._export_document_embeddings_async(flow)
|
||||
|
||||
|
|
@ -190,7 +190,7 @@ class BulkClient:
|
|||
except:
|
||||
pass
|
||||
|
||||
async def _export_document_embeddings_async(self, flow: str):
|
||||
async def _export_document_embeddings_async(self, flow: str) -> Iterator[Dict[str, Any]]:
|
||||
"""Async implementation of document embeddings export"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/export/document-embeddings"
|
||||
if self.token:
|
||||
|
|
@ -200,11 +200,11 @@ class BulkClient:
|
|||
async for raw_message in websocket:
|
||||
yield json.loads(raw_message)
|
||||
|
||||
def import_entity_contexts(self, flow: str, contexts: Iterator[Dict[str, Any]], **kwargs) -> None:
|
||||
def import_entity_contexts(self, flow: str, contexts: Iterator[Dict[str, Any]], **kwargs: Any) -> None:
|
||||
"""Bulk import entity contexts via WebSocket"""
|
||||
self._run_async(self._import_entity_contexts_async(flow, contexts))
|
||||
|
||||
async def _import_entity_contexts_async(self, flow: str, contexts: Iterator[Dict[str, Any]]):
|
||||
async def _import_entity_contexts_async(self, flow: str, contexts: Iterator[Dict[str, Any]]) -> None:
|
||||
"""Async implementation of entity contexts import"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/import/entity-contexts"
|
||||
if self.token:
|
||||
|
|
@ -214,7 +214,7 @@ class BulkClient:
|
|||
for context in contexts:
|
||||
await websocket.send(json.dumps(context))
|
||||
|
||||
def export_entity_contexts(self, flow: str, **kwargs) -> Iterator[Dict[str, Any]]:
|
||||
def export_entity_contexts(self, flow: str, **kwargs: Any) -> Iterator[Dict[str, Any]]:
|
||||
"""Bulk export entity contexts via WebSocket"""
|
||||
async_gen = self._export_entity_contexts_async(flow)
|
||||
|
||||
|
|
@ -240,7 +240,7 @@ class BulkClient:
|
|||
except:
|
||||
pass
|
||||
|
||||
async def _export_entity_contexts_async(self, flow: str):
|
||||
async def _export_entity_contexts_async(self, flow: str) -> Iterator[Dict[str, Any]]:
|
||||
"""Async implementation of entity contexts export"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/export/entity-contexts"
|
||||
if self.token:
|
||||
|
|
@ -250,11 +250,11 @@ class BulkClient:
|
|||
async for raw_message in websocket:
|
||||
yield json.loads(raw_message)
|
||||
|
||||
def import_objects(self, flow: str, objects: Iterator[Dict[str, Any]], **kwargs) -> None:
|
||||
def import_objects(self, flow: str, objects: Iterator[Dict[str, Any]], **kwargs: Any) -> None:
|
||||
"""Bulk import objects via WebSocket"""
|
||||
self._run_async(self._import_objects_async(flow, objects))
|
||||
|
||||
async def _import_objects_async(self, flow: str, objects: Iterator[Dict[str, Any]]):
|
||||
async def _import_objects_async(self, flow: str, objects: Iterator[Dict[str, Any]]) -> None:
|
||||
"""Async implementation of objects import"""
|
||||
ws_url = f"{self.url}/api/v1/flow/{flow}/import/objects"
|
||||
if self.token:
|
||||
|
|
@ -264,7 +264,7 @@ class BulkClient:
|
|||
for obj in objects:
|
||||
await websocket.send(json.dumps(obj))
|
||||
|
||||
def close(self):
|
||||
def close(self) -> None:
|
||||
"""Close connections"""
|
||||
# Cleanup handled by context managers
|
||||
pass
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue