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

Fix embeddings integration

Browse source
This commit is contained in:
Cyber MacGeddon 2026-03-08 19:35:38 +00:00
parent 0a2ce47a88
commit ea95effbf2
5 changed files with 21 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

3
trustgraph-flow/trustgraph/agent/react/tools.py
View file

@ -154,7 +154,8 @@ class RowEmbeddingsQueryImpl:
logger.debug("Getting embeddings for row query...") logger.debug("Getting embeddings for row query...")
query_text = arguments.get("query") query_text = arguments.get("query")
vectors = await embeddings_client.embed(query_text) all_vectors = await embeddings_client.embed([query_text])
vectors = all_vectors[0] if all_vectors else []
# Now query row embeddings # Now query row embeddings
client = self.context("row-embeddings-query-request") client = self.context("row-embeddings-query-request")

4
trustgraph-flow/trustgraph/extract/kg/ontology/extract.py
View file

@ -148,8 +148,8 @@ class Processor(FlowProcessor):
# Detect embedding dimension by embedding a test string # Detect embedding dimension by embedding a test string
logger.info("Detecting embedding dimension from embeddings service...") logger.info("Detecting embedding dimension from embeddings service...")
test_embedding_response = await embeddings_client.embed("test") test_embedding_response = await embeddings_client.embed(["test"])
test_embedding = test_embedding_response[0] # Extract from [[vector]] test_embedding = test_embedding_response[0][0] # Extract first vector from first text
dimension = len(test_embedding) dimension = len(test_embedding)
logger.info(f"Detected embedding dimension: {dimension}") logger.info(f"Detected embedding dimension: {dimension}")

25
trustgraph-flow/trustgraph/extract/kg/ontology/ontology_embedder.py
View file

@ -153,13 +153,11 @@ class OntologyEmbedder:
# Get embeddings for batch # Get embeddings for batch
texts = [elem['text'] for elem in batch] texts = [elem['text'] for elem in batch]
try: try:
# Call embedding service for each text # Single batch embedding call
# Note: embed() returns 2D array [[vector]], so extract first element embeddings_response = await self.embedding_service.embed(texts)
embedding_tasks = [self.embedding_service.embed(text) for text in texts]
embeddings_responses = await asyncio.gather(*embedding_tasks)
# Extract vectors from responses (each is [[vector]]) # Extract first vector from each text's vector set
embeddings_list = [resp[0] for resp in embeddings_responses] embeddings_list = [resp[0] for resp in embeddings_response]
# Convert to numpy array # Convert to numpy array
embeddings = np.array(embeddings_list) embeddings = np.array(embeddings_list)
@ -218,9 +216,9 @@ class OntologyEmbedder:
return None return None
try: try:
# embed() returns 2D array [[vector]], extract first element # embed() with single text, extract first vector from first text
embedding_response = await self.embedding_service.embed(text) embedding_response = await self.embedding_service.embed([text])
return np.array(embedding_response[0]) return np.array(embedding_response[0][0])
except Exception as e: except Exception as e:
logger.error(f"Failed to embed text: {e}") logger.error(f"Failed to embed text: {e}")
return None return None
@ -239,11 +237,10 @@ class OntologyEmbedder:
return None return None
try: try:
# Call embed() for each text (returns [[vector]] per call) # Single batch embedding call
embedding_tasks = [self.embedding_service.embed(text) for text in texts] embeddings_response = await self.embedding_service.embed(texts)
embeddings_responses = await asyncio.gather(*embedding_tasks) # Extract first vector from each text's vector set
# Extract first vector from each response embeddings_list = [resp[0] for resp in embeddings_response]
embeddings_list = [resp[0] for resp in embeddings_responses]
return np.array(embeddings_list) return np.array(embeddings_list)
except Exception as e: except Exception as e:
logger.error(f"Failed to embed texts: {e}") logger.error(f"Failed to embed texts: {e}")

5
trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py
View file

@ -24,12 +24,13 @@ class Query:
if self.verbose: if self.verbose:
logger.debug("Computing embeddings...") logger.debug("Computing embeddings...")
qembeds = await self.rag.embeddings_client.embed(query) qembeds = await self.rag.embeddings_client.embed([query])
if self.verbose: if self.verbose:
logger.debug("Embeddings computed") logger.debug("Embeddings computed")
return qembeds # Return the vector set for the first (only) text
return qembeds[0] if qembeds else []
async def get_docs(self, query): async def get_docs(self, query):

5
trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py
View file

@ -72,12 +72,13 @@ class Query:
if self.verbose: if self.verbose:
logger.debug("Computing embeddings...") logger.debug("Computing embeddings...")
qembeds = await self.rag.embeddings_client.embed(query) qembeds = await self.rag.embeddings_client.embed([query])
if self.verbose: if self.verbose:
logger.debug("Done.") logger.debug("Done.")
return qembeds # Return the vector set for the first (only) text
return qembeds[0] if qembeds else []
async def get_entities(self, query): async def get_entities(self, query):