mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-09 13:22:10 +02:00
Python API implements streaming interfaces (#577)
* Tech spec * Python CLI utilities updated to use the API including streaming features * Added type safety to Python API * Completed missing auth token support in CLI
This commit is contained in:
cybermaggedon
GitHub
parent
b957004db9
commit
01aeede78b
53 changed files with 4489 additions and 715 deletions
|
|
@ -2,18 +2,17 @@
|
|||
Loads triples and entity contexts into the knowledge graph.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import argparse
|
||||
import os
|
||||
import time
|
||||
import rdflib
|
||||
import json
|
||||
from websockets.asyncio.client import connect
|
||||
from typing import List, Dict, Any
|
||||
from typing import Iterator, Tuple
|
||||
|
||||
from trustgraph.api import Api, Triple
|
||||
from trustgraph.log_level import LogLevel
|
||||
|
||||
default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
|
||||
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
|
||||
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
|
||||
default_user = 'trustgraph'
|
||||
default_collection = 'default'
|
||||
|
||||
|
|
@ -26,108 +25,114 @@ class KnowledgeLoader:
|
|||
user,
|
||||
collection,
|
||||
document_id,
|
||||
url = default_url,
|
||||
url=default_url,
|
||||
token=None,
|
||||
):
|
||||
|
||||
if not url.endswith("/"):
|
||||
url += "/"
|
||||
|
||||
self.triples_url = url + f"api/v1/flow/{flow}/import/triples"
|
||||
self.entity_contexts_url = url + f"api/v1/flow/{flow}/import/entity-contexts"
|
||||
|
||||
self.files = files
|
||||
self.flow = flow
|
||||
self.user = user
|
||||
self.collection = collection
|
||||
self.document_id = document_id
|
||||
self.url = url
|
||||
self.token = token
|
||||
|
||||
async def run(self):
|
||||
|
||||
try:
|
||||
# Load triples first
|
||||
async with connect(self.triples_url) as ws:
|
||||
for file in self.files:
|
||||
await self.load_triples(file, ws)
|
||||
|
||||
# Then load entity contexts
|
||||
async with connect(self.entity_contexts_url) as ws:
|
||||
for file in self.files:
|
||||
await self.load_entity_contexts(file, ws)
|
||||
|
||||
except Exception as e:
|
||||
print(e, flush=True)
|
||||
|
||||
async def load_triples(self, file, ws):
|
||||
def load_triples_from_file(self, file) -> Iterator[Triple]:
|
||||
"""Generator that yields Triple objects from a Turtle file"""
|
||||
|
||||
g = rdflib.Graph()
|
||||
g.parse(file, format="turtle")
|
||||
|
||||
def Value(value, is_uri):
|
||||
return { "v": value, "e": is_uri }
|
||||
|
||||
for e in g:
|
||||
s = Value(value=str(e[0]), is_uri=True)
|
||||
p = Value(value=str(e[1]), is_uri=True)
|
||||
if type(e[2]) == rdflib.term.URIRef:
|
||||
o = Value(value=str(e[2]), is_uri=True)
|
||||
# Extract subject, predicate, object
|
||||
s_value = str(e[0])
|
||||
p_value = str(e[1])
|
||||
|
||||
# Check if object is a URI or literal
|
||||
if isinstance(e[2], rdflib.term.URIRef):
|
||||
o_value = str(e[2])
|
||||
o_is_uri = True
|
||||
else:
|
||||
o = Value(value=str(e[2]), is_uri=False)
|
||||
o_value = str(e[2])
|
||||
o_is_uri = False
|
||||
|
||||
req = {
|
||||
"metadata": {
|
||||
"id": self.document_id,
|
||||
"metadata": [],
|
||||
"user": self.user,
|
||||
"collection": self.collection
|
||||
},
|
||||
"triples": [
|
||||
{
|
||||
"s": s,
|
||||
"p": p,
|
||||
"o": o,
|
||||
}
|
||||
]
|
||||
}
|
||||
# Create Triple object
|
||||
# Note: The Triple dataclass has 's', 'p', 'o' fields as strings
|
||||
# The API will handle the metadata wrapping
|
||||
yield Triple(s=s_value, p=p_value, o=o_value)
|
||||
|
||||
await ws.send(json.dumps(req))
|
||||
|
||||
async def load_entity_contexts(self, file, ws):
|
||||
"""
|
||||
Load entity contexts by extracting entities from the RDF graph
|
||||
and generating contextual descriptions based on their relationships.
|
||||
"""
|
||||
def load_entity_contexts_from_file(self, file) -> Iterator[Tuple[str, str]]:
|
||||
"""Generator that yields (entity, context) tuples from a Turtle file"""
|
||||
|
||||
g = rdflib.Graph()
|
||||
g.parse(file, format="turtle")
|
||||
|
||||
for s, p, o in g:
|
||||
# If object is a URI, do nothing
|
||||
# If object is a URI, skip (we only want literal contexts)
|
||||
if isinstance(o, rdflib.term.URIRef):
|
||||
continue
|
||||
|
||||
# If object is a literal, create entity context for subject with literal as context
|
||||
|
||||
# If object is a literal, create entity context for subject
|
||||
s_str = str(s)
|
||||
o_str = str(o)
|
||||
|
||||
req = {
|
||||
"metadata": {
|
||||
"id": self.document_id,
|
||||
"metadata": [],
|
||||
"user": self.user,
|
||||
"collection": self.collection
|
||||
},
|
||||
"entities": [
|
||||
{
|
||||
"entity": {
|
||||
"v": s_str,
|
||||
"e": True
|
||||
},
|
||||
"context": o_str
|
||||
|
||||
yield (s_str, o_str)
|
||||
|
||||
def run(self):
|
||||
"""Load triples and entity contexts using Python API"""
|
||||
|
||||
try:
|
||||
# Create API client
|
||||
api = Api(url=self.url, token=self.token)
|
||||
bulk = api.bulk()
|
||||
|
||||
# Load triples from all files
|
||||
print("Loading triples...")
|
||||
for file in self.files:
|
||||
print(f" Processing {file}...")
|
||||
triples = self.load_triples_from_file(file)
|
||||
|
||||
bulk.import_triples(
|
||||
flow=self.flow,
|
||||
triples=triples,
|
||||
metadata={
|
||||
"id": self.document_id,
|
||||
"metadata": [],
|
||||
"user": self.user,
|
||||
"collection": self.collection
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
await ws.send(json.dumps(req))
|
||||
print("Triples loaded.")
|
||||
|
||||
# Load entity contexts from all files
|
||||
print("Loading entity contexts...")
|
||||
for file in self.files:
|
||||
print(f" Processing {file}...")
|
||||
|
||||
# Convert tuples to the format expected by import_entity_contexts
|
||||
def entity_context_generator():
|
||||
for entity, context in self.load_entity_contexts_from_file(file):
|
||||
yield {
|
||||
"entity": {"v": entity, "e": True},
|
||||
"context": context
|
||||
}
|
||||
|
||||
bulk.import_entity_contexts(
|
||||
flow=self.flow,
|
||||
entities=entity_context_generator(),
|
||||
metadata={
|
||||
"id": self.document_id,
|
||||
"metadata": [],
|
||||
"user": self.user,
|
||||
"collection": self.collection
|
||||
}
|
||||
)
|
||||
|
||||
print("Entity contexts loaded.")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error: {e}", flush=True)
|
||||
raise
|
||||
|
||||
def main():
|
||||
|
||||
|
|
@ -142,6 +147,12 @@ def main():
|
|||
help=f'API URL (default: {default_url})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-t', '--token',
|
||||
default=default_token,
|
||||
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-i', '--document-id',
|
||||
required=True,
|
||||
|
|
@ -166,7 +177,6 @@ def main():
|
|||
help=f'Collection ID (default: {default_collection})'
|
||||
)
|
||||
|
||||
|
||||
parser.add_argument(
|
||||
'files', nargs='+',
|
||||
help=f'Turtle files to load'
|
||||
|
|
@ -178,15 +188,16 @@ def main():
|
|||
|
||||
try:
|
||||
loader = KnowledgeLoader(
|
||||
document_id = args.document_id,
|
||||
url = args.api_url,
|
||||
flow = args.flow_id,
|
||||
files = args.files,
|
||||
user = args.user,
|
||||
collection = args.collection,
|
||||
document_id=args.document_id,
|
||||
url=args.api_url,
|
||||
token=args.token,
|
||||
flow=args.flow_id,
|
||||
files=args.files,
|
||||
user=args.user,
|
||||
collection=args.collection,
|
||||
)
|
||||
|
||||
asyncio.run(loader.run())
|
||||
loader.run()
|
||||
|
||||
print("Triples and entity contexts loaded.")
|
||||
break
|
||||
|
|
@ -199,4 +210,4 @@ def main():
|
|||
time.sleep(10)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue