apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-16 16:51:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Fixed DE queue schema

Browse source
This commit is contained in:
Cyber MacGeddon 2025-01-04 19:38:57 +00:00
parent 18e320eef7
commit 6e0bb36d75
5 changed files with 30 additions and 43 deletions
Download patch file Download diff file Expand all files Collapse all files

39
trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py
View file

@ -5,9 +5,9 @@ chunk of text. Input is chunk of text plus metadata.
Output is chunk plus embedding.
"""
from ... schema import EntityContexts, EntityEmbeddings, GraphEmbeddings
from ... schema import entity_contexts_ingest_queue
from ... schema import graph_embeddings_store_queue
from ... schema import Chunk, ChunkEmbeddings, DocumentEmbeddings
from ... schema import chunk_ingest_queue
from ... schema import document_embeddings_store_queue
from ... schema import embeddings_request_queue, embeddings_response_queue
from ... clients.embeddings_client import EmbeddingsClient
from ... log_level import LogLevel
@ -15,8 +15,8 @@ from ... base import ConsumerProducer
module = ".".join(__name__.split(".")[1:-1])
default_input_queue = entity_contexts_ingest_queue
default_output_queue = graph_embeddings_store_queue
default_input_queue = chunk_ingest_queue
default_output_queue = document_embeddings_store_queue
default_subscriber = module
class Processor(ConsumerProducer):
@ -40,8 +40,8 @@ class Processor(ConsumerProducer):
"embeddings_request_queue": emb_request_queue,
"embeddings_response_queue": emb_response_queue,
"subscriber": subscriber,
"input_schema": EntityContexts,
"output_schema": GraphEmbeddings,
"input_schema": Chunk,
"output_schema": DocumentEmbeddings,
}
)
@ -52,34 +52,25 @@ class Processor(ConsumerProducer):
subscriber=module + "-emb",
)
def emit(self, rec, vectors):
r = GraphEmbeddings(metadata=metadata, chunk=chunk, vectors=vectors)
self.producer.send(r)
def handle(self, msg):
v = msg.value()
print(f"Indexing {v.metadata.id}...", flush=True)
entities = []
try:
for entity in v.entities:
vectors = self.embeddings.request(v.chunk)
vectors = self.embeddings.request(entity.context)
entities.append(
EntityEmbeddings(
entity=entity.entity,
vectors=vectors
)
embeds = [
ChunkEmbeddings(
chunk=v.chunk,
vectors=vectors,
)
]
r = GraphEmbeddings(
r = DocumentEmbeddings(
metadata=v.metadata,
entities=entities,
chunks=embeds,
)
self.producer.send(r)

5
trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py
View file

@ -52,11 +52,6 @@ class Processor(ConsumerProducer):
subscriber=module + "-emb",
)
def emit(self, rec, vectors):
r = GraphEmbeddings(metadata=metadata, chunk=chunk, vectors=vectors)
self.producer.send(r)
def handle(self, msg):
v = msg.value()

11
trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py
View file

@ -3,15 +3,16 @@
Accepts entity/vector pairs and writes them to a Milvus store.
"""
from .... schema import ChunkEmbeddings
from .... schema import chunk_embeddings_ingest_queue
from .... log_level import LogLevel
from .... direct.milvus_doc_embeddings import DocVectors
from .... schema import DocumentEmbeddings
from .... schema import document_embeddings_store_queue
from .... log_level import LogLevel
from .... base import Consumer
module = ".".join(__name__.split(".")[1:-1])
default_input_queue = chunk_embeddings_ingest_queue
default_input_queue = document_embeddings_store_queue
default_subscriber = module
default_store_uri = 'http://localhost:19530'
@ -27,7 +28,7 @@ class Processor(Consumer):
**params | {
"input_queue": input_queue,
"subscriber": subscriber,
"input_schema": ChunkEmbeddings,
"input_schema": DocumentEmbeddings,
"store_uri": store_uri,
}
)

8
trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py
View file

@ -11,14 +11,14 @@ import time
import uuid
import os
from .... schema import ChunkEmbeddings
from .... schema import chunk_embeddings_ingest_queue
from .... schema import DocumentEmbeddings
from .... schema import document_embeddings_store_queue
from .... log_level import LogLevel
from .... base import Consumer
module = ".".join(__name__.split(".")[1:-1])
default_input_queue = chunk_embeddings_ingest_queue
default_input_queue = document_embeddings_store_queue
default_subscriber = module
default_api_key = os.getenv("PINECONE_API_KEY", "not-specified")
default_cloud = "aws"
@ -54,7 +54,7 @@ class Processor(Consumer):
**params | {
"input_queue": input_queue,
"subscriber": subscriber,
"input_schema": ChunkEmbeddings,
"input_schema": DocumentEmbeddings,
"url": self.url,
}
)

10
trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py
View file

@ -8,14 +8,14 @@ from qdrant_client.models import PointStruct
from qdrant_client.models import Distance, VectorParams
import uuid
from .... schema import ChunkEmbeddings
from .... schema import chunk_embeddings_ingest_queue
from .... schema import DocumentEmbeddings
from .... schema import document_embeddings_store_queue
from .... log_level import LogLevel
from .... base import Consumer
module = ".".join(__name__.split(".")[1:-1])
default_input_queue = chunk_embeddings_ingest_queue
default_input_queue = document_embeddings_store_queue
default_subscriber = module
default_store_uri = 'http://localhost:6333'
@ -31,7 +31,7 @@ class Processor(Consumer):
**params | {
"input_queue": input_queue,
"subscriber": subscriber,
"input_schema": ChunkEmbeddings,
"input_schema": DocumentEmbeddings,
"store_uri": store_uri,
}
)
@ -64,7 +64,7 @@ class Processor(Consumer):
self.client.create_collection(
collection_name=collection,
vectors_config=VectorParams(
size=dim, distance=Distance.DOT
size=dim, distance=Distance.COSINE
),
)
except Exception as e: