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:30:39 +00:00
parent aca5ee5653
commit a7ae5a8823
9 changed files with 52 additions and 52 deletions
Download patch file Download diff file Expand all files Collapse all files

26
trustgraph-base/trustgraph/api/async_flow.py
View file

@ -613,8 +613,8 @@ class AsyncFlowInstance:
``` ```
""" """
# First convert text to embeddings vectors # First convert text to embeddings vectors
emb_result = await self.embeddings(text=text) emb_result = await self.embeddings(texts=[text])
vectors = emb_result.get("vectors", []) vectors = emb_result.get("vectors", [[]])[0]
request_data = { request_data = {
"vectors": vectors, "vectors": vectors,
@ -626,20 +626,20 @@ class AsyncFlowInstance:
return await self.request("graph-embeddings", request_data) return await self.request("graph-embeddings", request_data)
async def embeddings(self, text: str, **kwargs: Any): async def embeddings(self, texts: list, **kwargs: Any):
""" """
Generate embeddings for input text. Generate embeddings for input texts.
Converts text into a numerical vector representation using the flow's Converts texts into numerical vector representations using the flow's
configured embedding model. Useful for semantic search and similarity configured embedding model. Useful for semantic search and similarity
comparisons. comparisons.
Args: Args:
text: Input text to embed texts: List of input texts to embed
**kwargs: Additional service-specific parameters **kwargs: Additional service-specific parameters
Returns: Returns:
dict: Response containing embedding vector and metadata dict: Response containing embedding vectors
Example: Example:
```python ```python
@ -647,12 +647,12 @@ class AsyncFlowInstance:
flow = async_flow.id("default") flow = async_flow.id("default")
# Generate embeddings # Generate embeddings
result = await flow.embeddings(text="Sample text to embed") result = await flow.embeddings(texts=["Sample text to embed"])
vector = result.get("embedding") vectors = result.get("vectors")
print(f"Embedding dimension: {len(vector)}") print(f"Embedding dimension: {len(vectors[0][0])}")
``` ```
""" """
request_data = {"text": text} request_data = {"texts": texts}
request_data.update(kwargs) request_data.update(kwargs)
return await self.request("embeddings", request_data) return await self.request("embeddings", request_data)
@ -811,8 +811,8 @@ class AsyncFlowInstance:
``` ```
""" """
# First convert text to embeddings vectors # First convert text to embeddings vectors
emb_result = await self.embeddings(text=text) emb_result = await self.embeddings(texts=[text])
vectors = emb_result.get("vectors", []) vectors = emb_result.get("vectors", [[]])[0]
request_data = { request_data = {
"vectors": vectors, "vectors": vectors,

12
trustgraph-base/trustgraph/api/async_socket_client.py
View file

@ -283,8 +283,8 @@ class AsyncSocketFlowInstance:
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):
"""Query graph embeddings for semantic search""" """Query graph embeddings for semantic search"""
# First convert text to embeddings vectors # First convert text to embeddings vectors
emb_result = await self.embeddings(text=text) emb_result = await self.embeddings(texts=[text])
vectors = emb_result.get("vectors", []) vectors = emb_result.get("vectors", [[]])[0]
request = { request = {
"vectors": vectors, "vectors": vectors,
@ -296,9 +296,9 @@ class AsyncSocketFlowInstance:
return await self.client._send_request("graph-embeddings", self.flow_id, request) return await self.client._send_request("graph-embeddings", self.flow_id, request)
async def embeddings(self, text: str, **kwargs): async def embeddings(self, texts: list, **kwargs):
"""Generate text embeddings""" """Generate text embeddings"""
request = {"text": text} request = {"texts": texts}
request.update(kwargs) request.update(kwargs)
return await self.client._send_request("embeddings", self.flow_id, request) return await self.client._send_request("embeddings", self.flow_id, request)
@ -353,8 +353,8 @@ class AsyncSocketFlowInstance:
): ):
"""Query row embeddings for semantic search on structured data""" """Query row embeddings for semantic search on structured data"""
# First convert text to embeddings vectors # First convert text to embeddings vectors
emb_result = await self.embeddings(text=text) emb_result = await self.embeddings(texts=[text])
vectors = emb_result.get("vectors", []) vectors = emb_result.get("vectors", [[]])[0]
request = { request = {
"vectors": vectors, "vectors": vectors,

12
trustgraph-base/trustgraph/api/flow.py
View file

@ -603,8 +603,8 @@ class FlowInstance:
""" """
# First convert text to embeddings vectors # First convert text to embeddings vectors
emb_result = self.embeddings(text=text) emb_result = self.embeddings(texts=[text])
vectors = emb_result.get("vectors", []) vectors = emb_result.get("vectors", [[]])[0]
# Query graph embeddings for semantic search # Query graph embeddings for semantic search
input = { input = {
@ -649,8 +649,8 @@ class FlowInstance:
""" """
# First convert text to embeddings vectors # First convert text to embeddings vectors
emb_result = self.embeddings(text=text) emb_result = self.embeddings(texts=[text])
vectors = emb_result.get("vectors", []) vectors = emb_result.get("vectors", [[]])[0]
# Query document embeddings for semantic search # Query document embeddings for semantic search
input = { input = {
@ -1363,8 +1363,8 @@ class FlowInstance:
""" """
# First convert text to embeddings vectors # First convert text to embeddings vectors
emb_result = self.embeddings(text=text) emb_result = self.embeddings(texts=[text])
vectors = emb_result.get("vectors", []) vectors = emb_result.get("vectors", [[]])[0]
# Query row embeddings for semantic search # Query row embeddings for semantic search
input = { input = {

12
trustgraph-base/trustgraph/api/socket_client.py
View file

@ -650,8 +650,8 @@ class SocketFlowInstance:
``` ```
""" """
# First convert text to embeddings vectors # First convert text to embeddings vectors
emb_result = self.embeddings(text=text) emb_result = self.embeddings(texts=[text])
vectors = emb_result.get("vectors", []) vectors = emb_result.get("vectors", [[]])[0]
request = { request = {
"vectors": vectors, "vectors": vectors,
@ -699,8 +699,8 @@ class SocketFlowInstance:
``` ```
""" """
# First convert text to embeddings vectors # First convert text to embeddings vectors
emb_result = self.embeddings(text=text) emb_result = self.embeddings(texts=[text])
vectors = emb_result.get("vectors", []) vectors = emb_result.get("vectors", [[]])[0]
request = { request = {
"vectors": vectors, "vectors": vectors,
@ -937,8 +937,8 @@ class SocketFlowInstance:
``` ```
""" """
# First convert text to embeddings vectors # First convert text to embeddings vectors
emb_result = self.embeddings(text=text) emb_result = self.embeddings(texts=[text])
vectors = emb_result.get("vectors", []) vectors = emb_result.get("vectors", [[]])[0]
request = { request = {
"vectors": vectors, "vectors": vectors,

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):