mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-16 16:51:02 +02:00
Fix doc RAG processing, put graph-rag and doc-rag in appropriate component
files.
This commit is contained in:
Cyber MacGeddon
parent
906417c54c
commit
6c3cd607d1
11 changed files with 245 additions and 180 deletions
|
|
@ -39,5 +39,35 @@ local prompts = import "prompts/mixtral.jsonnet";
|
|||
|
||||
},
|
||||
|
||||
"document-embeddings" +: {
|
||||
|
||||
create:: function(engine)
|
||||
|
||||
local container =
|
||||
engine.container("document-embeddings")
|
||||
.with_image(images.trustgraph)
|
||||
.with_command([
|
||||
"document-embeddings",
|
||||
"-p",
|
||||
url.pulsar,
|
||||
])
|
||||
.with_limits("1.0", "512M")
|
||||
.with_reservations("0.5", "512M");
|
||||
|
||||
local containerSet = engine.containers(
|
||||
"document-embeddings", [ container ]
|
||||
);
|
||||
|
||||
local service =
|
||||
engine.internalService(containerSet)
|
||||
.with_port(8000, 8000, "metrics");
|
||||
|
||||
engine.resources([
|
||||
containerSet,
|
||||
service,
|
||||
])
|
||||
|
||||
},
|
||||
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -138,5 +138,35 @@ local url = import "values/url.jsonnet";
|
|||
|
||||
},
|
||||
|
||||
"graph-embeddings" +: {
|
||||
|
||||
create:: function(engine)
|
||||
|
||||
local container =
|
||||
engine.container("graph-embeddings")
|
||||
.with_image(images.trustgraph)
|
||||
.with_command([
|
||||
"graph-embeddings",
|
||||
"-p",
|
||||
url.pulsar,
|
||||
])
|
||||
.with_limits("1.0", "512M")
|
||||
.with_reservations("0.5", "512M");
|
||||
|
||||
local containerSet = engine.containers(
|
||||
"graph-embeddings", [ container ]
|
||||
);
|
||||
|
||||
local service =
|
||||
engine.internalService(containerSet)
|
||||
.with_port(8000, 8000, "metrics");
|
||||
|
||||
engine.resources([
|
||||
containerSet,
|
||||
service,
|
||||
])
|
||||
|
||||
},
|
||||
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -119,66 +119,6 @@ local prompt = import "prompt-template.jsonnet";
|
|||
|
||||
},
|
||||
|
||||
"graph-embeddings" +: {
|
||||
|
||||
create:: function(engine)
|
||||
|
||||
local container =
|
||||
engine.container("graph-embeddings")
|
||||
.with_image(images.trustgraph)
|
||||
.with_command([
|
||||
"graph-embeddings",
|
||||
"-p",
|
||||
url.pulsar,
|
||||
])
|
||||
.with_limits("1.0", "512M")
|
||||
.with_reservations("0.5", "512M");
|
||||
|
||||
local containerSet = engine.containers(
|
||||
"graph-embeddings", [ container ]
|
||||
);
|
||||
|
||||
local service =
|
||||
engine.internalService(containerSet)
|
||||
.with_port(8000, 8000, "metrics");
|
||||
|
||||
engine.resources([
|
||||
containerSet,
|
||||
service,
|
||||
])
|
||||
|
||||
},
|
||||
|
||||
"document-embeddings" +: {
|
||||
|
||||
create:: function(engine)
|
||||
|
||||
local container =
|
||||
engine.container("document-embeddings")
|
||||
.with_image(images.trustgraph)
|
||||
.with_command([
|
||||
"document-embeddings",
|
||||
"-p",
|
||||
url.pulsar,
|
||||
])
|
||||
.with_limits("1.0", "512M")
|
||||
.with_reservations("0.5", "512M");
|
||||
|
||||
local containerSet = engine.containers(
|
||||
"document-embeddings", [ container ]
|
||||
);
|
||||
|
||||
local service =
|
||||
engine.internalService(containerSet)
|
||||
.with_port(8000, 8000, "metrics");
|
||||
|
||||
engine.resources([
|
||||
containerSet,
|
||||
service,
|
||||
])
|
||||
|
||||
},
|
||||
|
||||
"metering" +: {
|
||||
|
||||
create:: function(engine)
|
||||
|
|
|
|||
|
|
@ -131,6 +131,35 @@ class Api:
|
|||
except:
|
||||
raise ProtocolException(f"Response not formatted correctly")
|
||||
|
||||
def document_rag(self, question):
|
||||
|
||||
# The input consists of a question
|
||||
input = {
|
||||
"query": question
|
||||
}
|
||||
|
||||
url = f"{self.url}document-rag"
|
||||
|
||||
# Invoke the API, input is passed as JSON
|
||||
resp = requests.post(url, json=input)
|
||||
|
||||
# Should be a 200 status code
|
||||
if resp.status_code != 200:
|
||||
raise ProtocolException(f"Status code {resp.status_code}")
|
||||
|
||||
try:
|
||||
# Parse the response as JSON
|
||||
object = resp.json()
|
||||
except:
|
||||
raise ProtocolException(f"Expected JSON response")
|
||||
|
||||
self.check_error(resp)
|
||||
|
||||
try:
|
||||
return object["response"]
|
||||
except:
|
||||
raise ProtocolException(f"Response not formatted correctly")
|
||||
|
||||
def embeddings(self, text):
|
||||
|
||||
# The input consists of a text block
|
||||
|
|
|
|||
|
|
@ -38,8 +38,12 @@ class DocumentEmbeddingsClient(BaseClient):
|
|||
output_schema=DocumentEmbeddingsResponse,
|
||||
)
|
||||
|
||||
def request(self, vectors, limit=10, timeout=300):
|
||||
def request(
|
||||
self, vectors, user="trustgraph", collection="default",
|
||||
limit=10, timeout=300
|
||||
):
|
||||
return self.call(
|
||||
user=user, collection=collection,
|
||||
vectors=vectors, limit=limit, timeout=timeout
|
||||
).documents
|
||||
|
||||
|
|
|
|||
|
|
@ -55,6 +55,8 @@ document_embeddings_store_queue = topic('document-embeddings-store')
|
|||
class DocumentEmbeddingsRequest(Record):
|
||||
vectors = Array(Array(Double()))
|
||||
limit = Integer()
|
||||
user = String()
|
||||
collection = String()
|
||||
|
||||
class DocumentEmbeddingsResponse(Record):
|
||||
error = Error()
|
||||
|
|
|
|||
|
|
@ -16,6 +16,44 @@ from . schema import document_embeddings_response_queue
|
|||
LABEL="http://www.w3.org/2000/01/rdf-schema#label"
|
||||
DEFINITION="http://www.w3.org/2004/02/skos/core#definition"
|
||||
|
||||
class Query:
|
||||
|
||||
def __init__(self, rag, user, collection, verbose):
|
||||
self.rag = rag
|
||||
self.user = user
|
||||
self.collection = collection
|
||||
self.verbose = verbose
|
||||
|
||||
def get_vector(self, query):
|
||||
|
||||
if self.verbose:
|
||||
print("Compute embeddings...", flush=True)
|
||||
|
||||
qembeds = self.rag.embeddings.request(query)
|
||||
|
||||
if self.verbose:
|
||||
print("Done.", flush=True)
|
||||
|
||||
return qembeds
|
||||
|
||||
def get_docs(self, query):
|
||||
|
||||
vectors = self.get_vector(query)
|
||||
|
||||
if self.verbose:
|
||||
print("Get entities...", flush=True)
|
||||
|
||||
docs = self.rag.de_client.request(
|
||||
vectors, limit=self.rag.doc_limit
|
||||
)
|
||||
|
||||
if self.verbose:
|
||||
print("Docs:", flush=True)
|
||||
for doc in docs:
|
||||
print(doc, flush=True)
|
||||
|
||||
return docs
|
||||
|
||||
class DocumentRag:
|
||||
|
||||
def __init__(
|
||||
|
|
@ -55,7 +93,7 @@ class DocumentRag:
|
|||
print("Initialising...", flush=True)
|
||||
|
||||
# FIXME: Configurable
|
||||
self.entity_limit = 20
|
||||
self.doc_limit = 20
|
||||
|
||||
self.de_client = DocumentEmbeddingsClient(
|
||||
pulsar_host=pulsar_host,
|
||||
|
|
@ -81,42 +119,16 @@ class DocumentRag:
|
|||
if self.verbose:
|
||||
print("Initialised", flush=True)
|
||||
|
||||
def get_vector(self, query):
|
||||
|
||||
if self.verbose:
|
||||
print("Compute embeddings...", flush=True)
|
||||
|
||||
qembeds = self.embeddings.request(query)
|
||||
|
||||
if self.verbose:
|
||||
print("Done.", flush=True)
|
||||
|
||||
return qembeds
|
||||
|
||||
def get_docs(self, query):
|
||||
|
||||
vectors = self.get_vector(query)
|
||||
|
||||
if self.verbose:
|
||||
print("Get entities...", flush=True)
|
||||
|
||||
docs = self.de_client.request(
|
||||
vectors, self.entity_limit
|
||||
)
|
||||
|
||||
if self.verbose:
|
||||
print("Docs:", flush=True)
|
||||
for doc in docs:
|
||||
print(doc, flush=True)
|
||||
|
||||
return docs
|
||||
|
||||
def query(self, query):
|
||||
def query(self, query, user="trustgraph", collection="default"):
|
||||
|
||||
if self.verbose:
|
||||
print("Construct prompt...", flush=True)
|
||||
|
||||
docs = self.get_docs(query)
|
||||
q = Query(
|
||||
rag=self, user=user, collection=collection, verbose=self.verbose
|
||||
)
|
||||
|
||||
docs = q.get_docs(query)
|
||||
|
||||
if self.verbose:
|
||||
print("Invoke LLM...", flush=True)
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ from . subscriber import Subscriber
|
|||
from . text_completion import TextCompletionRequestor
|
||||
from . prompt import PromptRequestor
|
||||
from . graph_rag import GraphRagRequestor
|
||||
from . document_rag import DocumentRagRequestor
|
||||
from . triples_query import TriplesQueryRequestor
|
||||
from . graph_embeddings_query import GraphEmbeddingsQueryRequestor
|
||||
from . embeddings import EmbeddingsRequestor
|
||||
|
|
@ -91,6 +92,10 @@ class Api:
|
|||
pulsar_host=self.pulsar_host, timeout=self.timeout,
|
||||
auth = self.auth,
|
||||
),
|
||||
"document-rag": DocumentRagRequestor(
|
||||
pulsar_host=self.pulsar_host, timeout=self.timeout,
|
||||
auth = self.auth,
|
||||
),
|
||||
"triples-query": TriplesQueryRequestor(
|
||||
pulsar_host=self.pulsar_host, timeout=self.timeout,
|
||||
auth = self.auth,
|
||||
|
|
@ -140,6 +145,10 @@ class Api:
|
|||
endpoint_path = "/api/v1/graph-rag", auth=self.auth,
|
||||
requestor = self.services["graph-rag"],
|
||||
),
|
||||
ServiceEndpoint(
|
||||
endpoint_path = "/api/v1/document-rag", auth=self.auth,
|
||||
requestor = self.services["document-rag"],
|
||||
),
|
||||
ServiceEndpoint(
|
||||
endpoint_path = "/api/v1/triples-query", auth=self.auth,
|
||||
requestor = self.services["triples-query"],
|
||||
|
|
|
|||
|
|
@ -39,11 +39,16 @@ class Processor(Consumer):
|
|||
|
||||
v = msg.value()
|
||||
|
||||
chunk = v.chunk.decode("utf-8")
|
||||
for emb in v.chunks:
|
||||
|
||||
if v.chunk != "" and v.chunk is not None:
|
||||
for vec in v.vectors:
|
||||
self.vecstore.insert(vec, chunk)
|
||||
chunk = emb.chunk.decode("utf-8")
|
||||
if chunk == "" or chunk is None: continue
|
||||
|
||||
for vec in emb.vectors:
|
||||
|
||||
if chunk != "" and v.chunk is not None:
|
||||
for vec in v.vectors:
|
||||
self.vecstore.insert(vec, chunk)
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
|
|
|||
|
|
@ -65,71 +65,74 @@ class Processor(Consumer):
|
|||
|
||||
v = msg.value()
|
||||
|
||||
chunk = v.chunk.decode("utf-8")
|
||||
for emb in v.chunks:
|
||||
|
||||
if chunk == "": return
|
||||
chunk = emb.chunk.decode("utf-8")
|
||||
if chunk == "" or chunk is None: continue
|
||||
|
||||
for vec in v.vectors:
|
||||
for vec in emb.vectors:
|
||||
|
||||
dim = len(vec)
|
||||
collection = (
|
||||
"d-" + v.metadata.user + "-" + str(dim)
|
||||
)
|
||||
for vec in v.vectors:
|
||||
|
||||
if index_name != self.last_index_name:
|
||||
dim = len(vec)
|
||||
collection = (
|
||||
"d-" + v.metadata.user + "-" + str(dim)
|
||||
)
|
||||
|
||||
if not self.pinecone.has_index(index_name):
|
||||
if index_name != self.last_index_name:
|
||||
|
||||
try:
|
||||
if not self.pinecone.has_index(index_name):
|
||||
|
||||
self.pinecone.create_index(
|
||||
name = index_name,
|
||||
dimension = dim,
|
||||
metric = "cosine",
|
||||
spec = ServerlessSpec(
|
||||
cloud = self.cloud,
|
||||
region = self.region,
|
||||
)
|
||||
)
|
||||
try:
|
||||
|
||||
for i in range(0, 1000):
|
||||
self.pinecone.create_index(
|
||||
name = index_name,
|
||||
dimension = dim,
|
||||
metric = "cosine",
|
||||
spec = ServerlessSpec(
|
||||
cloud = self.cloud,
|
||||
region = self.region,
|
||||
)
|
||||
)
|
||||
|
||||
if self.pinecone.describe_index(
|
||||
index_name
|
||||
).status["ready"]:
|
||||
break
|
||||
for i in range(0, 1000):
|
||||
|
||||
time.sleep(1)
|
||||
if self.pinecone.describe_index(
|
||||
index_name
|
||||
).status["ready"]:
|
||||
break
|
||||
|
||||
if not self.pinecone.describe_index(
|
||||
index_name
|
||||
).status["ready"]:
|
||||
raise RuntimeError(
|
||||
"Gave up waiting for index creation"
|
||||
)
|
||||
time.sleep(1)
|
||||
|
||||
except Exception as e:
|
||||
print("Pinecone index creation failed")
|
||||
raise e
|
||||
if not self.pinecone.describe_index(
|
||||
index_name
|
||||
).status["ready"]:
|
||||
raise RuntimeError(
|
||||
"Gave up waiting for index creation"
|
||||
)
|
||||
|
||||
print(f"Index {index_name} created", flush=True)
|
||||
except Exception as e:
|
||||
print("Pinecone index creation failed")
|
||||
raise e
|
||||
|
||||
self.last_index_name = index_name
|
||||
print(f"Index {index_name} created", flush=True)
|
||||
|
||||
index = self.pinecone.Index(index_name)
|
||||
self.last_index_name = index_name
|
||||
|
||||
records = [
|
||||
{
|
||||
"id": id,
|
||||
"values": vec,
|
||||
"metadata": { "doc": chunk },
|
||||
}
|
||||
]
|
||||
index = self.pinecone.Index(index_name)
|
||||
|
||||
index.upsert(
|
||||
vectors = records,
|
||||
namespace = v.metadata.collection,
|
||||
)
|
||||
records = [
|
||||
{
|
||||
"id": id,
|
||||
"values": vec,
|
||||
"metadata": { "doc": chunk },
|
||||
}
|
||||
]
|
||||
|
||||
index.upsert(
|
||||
vectors = records,
|
||||
namespace = v.metadata.collection,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
|
|
|||
|
|
@ -44,47 +44,48 @@ class Processor(Consumer):
|
|||
|
||||
v = msg.value()
|
||||
|
||||
chunk = v.chunk.decode("utf-8")
|
||||
for emb in v.chunks:
|
||||
|
||||
if chunk == "": return
|
||||
chunk = emb.chunk.decode("utf-8")
|
||||
if chunk == "": return
|
||||
|
||||
for vec in v.vectors:
|
||||
for vec in emb.vectors:
|
||||
|
||||
dim = len(vec)
|
||||
collection = (
|
||||
"d_" + v.metadata.user + "_" + v.metadata.collection + "_" +
|
||||
str(dim)
|
||||
)
|
||||
dim = len(vec)
|
||||
collection = (
|
||||
"d_" + v.metadata.user + "_" + v.metadata.collection + "_" +
|
||||
str(dim)
|
||||
)
|
||||
|
||||
if collection != self.last_collection:
|
||||
if collection != self.last_collection:
|
||||
|
||||
if not self.client.collection_exists(collection):
|
||||
if not self.client.collection_exists(collection):
|
||||
|
||||
try:
|
||||
self.client.create_collection(
|
||||
collection_name=collection,
|
||||
vectors_config=VectorParams(
|
||||
size=dim, distance=Distance.COSINE
|
||||
),
|
||||
try:
|
||||
self.client.create_collection(
|
||||
collection_name=collection,
|
||||
vectors_config=VectorParams(
|
||||
size=dim, distance=Distance.COSINE
|
||||
),
|
||||
)
|
||||
except Exception as e:
|
||||
print("Qdrant collection creation failed")
|
||||
raise e
|
||||
|
||||
self.last_collection = collection
|
||||
|
||||
self.client.upsert(
|
||||
collection_name=collection,
|
||||
points=[
|
||||
PointStruct(
|
||||
id=str(uuid.uuid4()),
|
||||
vector=vec,
|
||||
payload={
|
||||
"doc": chunk,
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
print("Qdrant collection creation failed")
|
||||
raise e
|
||||
|
||||
self.last_collection = collection
|
||||
|
||||
self.client.upsert(
|
||||
collection_name=collection,
|
||||
points=[
|
||||
PointStruct(
|
||||
id=str(uuid.uuid4()),
|
||||
vector=vec,
|
||||
payload={
|
||||
"doc": chunk,
|
||||
}
|
||||
)
|
||||
]
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue