mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-22 19:51:02 +02:00
Doc RAG working
This commit is contained in:
Cyber MacGeddon
parent
6be7b30633
commit
c9913297d2
14 changed files with 157 additions and 237 deletions
|
|
@ -7,9 +7,7 @@ as text as separate output objects.
|
|||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
from prometheus_client import Histogram
|
||||
|
||||
from ... schema import TextDocument, Chunk, Metadata
|
||||
from ... schema import text_ingest_queue, chunk_ingest_queue
|
||||
from ... log_level import LogLevel
|
||||
from ... schema import TextDocument, Chunk
|
||||
from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
|
||||
|
||||
default_ident = "chunker"
|
||||
|
|
|
|||
|
|
@ -7,9 +7,7 @@ as text as separate output objects.
|
|||
from langchain_text_splitters import TokenTextSplitter
|
||||
from prometheus_client import Histogram
|
||||
|
||||
from ... schema import TextDocument, Chunk, Metadata
|
||||
from ... schema import text_ingest_queue, chunk_ingest_queue
|
||||
from ... log_level import LogLevel
|
||||
from ... schema import TextDocument, Chunk
|
||||
from ... base import FlowProcessor
|
||||
|
||||
default_ident = "chunker"
|
||||
|
|
|
|||
|
|
@ -9,7 +9,6 @@ import base64
|
|||
from langchain_community.document_loaders import PyPDFLoader
|
||||
|
||||
from ... schema import Document, TextDocument, Metadata
|
||||
from ... schema import document_ingest_queue, text_ingest_queue
|
||||
from ... log_level import LogLevel
|
||||
from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
|
||||
|
||||
|
|
|
|||
|
|
@ -34,31 +34,24 @@ class Processor(DocumentEmbeddingsQueryService):
|
|||
|
||||
self.qdrant = QdrantClient(url=store_uri, api_key=api_key)
|
||||
|
||||
async def handle(self, msg):
|
||||
async def query_document_embeddings(self, msg):
|
||||
|
||||
try:
|
||||
|
||||
v = msg.value()
|
||||
|
||||
# Sender-produced ID
|
||||
id = msg.properties()["id"]
|
||||
|
||||
print(f"Handling input {id}...", flush=True)
|
||||
|
||||
chunks = []
|
||||
|
||||
for vec in v.vectors:
|
||||
for vec in msg.vectors:
|
||||
|
||||
dim = len(vec)
|
||||
collection = (
|
||||
"d_" + v.user + "_" + v.collection + "_" +
|
||||
"d_" + msg.user + "_" + msg.collection + "_" +
|
||||
str(dim)
|
||||
)
|
||||
|
||||
search_result = self.qdrant.query_points(
|
||||
collection_name=collection,
|
||||
query=vec,
|
||||
limit=v.limit,
|
||||
limit=msg.limit,
|
||||
with_payload=True,
|
||||
).points
|
||||
|
||||
|
|
|
|||
|
|
@ -1,20 +1,7 @@
|
|||
|
||||
from . clients.document_embeddings_client import DocumentEmbeddingsClient
|
||||
from . clients.triples_query_client import TriplesQueryClient
|
||||
from . clients.embeddings_client import EmbeddingsClient
|
||||
from . clients.prompt_client import PromptClient
|
||||
|
||||
from . schema import DocumentEmbeddingsRequest, DocumentEmbeddingsResponse
|
||||
from . schema import TriplesQueryRequest, TriplesQueryResponse
|
||||
from . schema import prompt_request_queue
|
||||
from . schema import prompt_response_queue
|
||||
from . schema import embeddings_request_queue
|
||||
from . schema import embeddings_response_queue
|
||||
from . schema import document_embeddings_request_queue
|
||||
from . schema import document_embeddings_response_queue
|
||||
import asyncio
|
||||
|
||||
LABEL="http://www.w3.org/2000/01/rdf-schema#label"
|
||||
DEFINITION="http://www.w3.org/2004/02/skos/core#definition"
|
||||
|
||||
class Query:
|
||||
|
||||
|
|
@ -28,27 +15,28 @@ class Query:
|
|||
self.verbose = verbose
|
||||
self.doc_limit = doc_limit
|
||||
|
||||
def get_vector(self, query):
|
||||
async def get_vector(self, query):
|
||||
|
||||
if self.verbose:
|
||||
print("Compute embeddings...", flush=True)
|
||||
|
||||
qembeds = self.rag.embeddings.request(query)
|
||||
qembeds = await self.rag.embeddings_client.embed(query)
|
||||
|
||||
if self.verbose:
|
||||
print("Done.", flush=True)
|
||||
|
||||
return qembeds
|
||||
|
||||
def get_docs(self, query):
|
||||
async def get_docs(self, query):
|
||||
|
||||
vectors = self.get_vector(query)
|
||||
vectors = await self.get_vector(query)
|
||||
|
||||
if self.verbose:
|
||||
print("Get entities...", flush=True)
|
||||
print("Get docs...", flush=True)
|
||||
|
||||
docs = self.rag.de_client.request(
|
||||
vectors, limit=self.doc_limit
|
||||
docs = await self.rag.doc_embeddings_client.query(
|
||||
vectors, limit=self.doc_limit,
|
||||
user=self.user, collection=self.collection,
|
||||
)
|
||||
|
||||
if self.verbose:
|
||||
|
|
@ -61,70 +49,20 @@ class Query:
|
|||
class DocumentRag:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pulsar_host="pulsar://pulsar:6650",
|
||||
pulsar_api_key=None,
|
||||
pr_request_queue=None,
|
||||
pr_response_queue=None,
|
||||
emb_request_queue=None,
|
||||
emb_response_queue=None,
|
||||
de_request_queue=None,
|
||||
de_response_queue=None,
|
||||
self, prompt_client, embeddings_client, doc_embeddings_client,
|
||||
verbose=False,
|
||||
module="test",
|
||||
):
|
||||
|
||||
self.verbose=verbose
|
||||
self.verbose = verbose
|
||||
|
||||
if pr_request_queue is None:
|
||||
pr_request_queue = prompt_request_queue
|
||||
|
||||
if pr_response_queue is None:
|
||||
pr_response_queue = prompt_response_queue
|
||||
|
||||
if emb_request_queue is None:
|
||||
emb_request_queue = embeddings_request_queue
|
||||
|
||||
if emb_response_queue is None:
|
||||
emb_response_queue = embeddings_response_queue
|
||||
|
||||
if de_request_queue is None:
|
||||
de_request_queue = document_embeddings_request_queue
|
||||
|
||||
if de_response_queue is None:
|
||||
de_response_queue = document_embeddings_response_queue
|
||||
|
||||
if self.verbose:
|
||||
print("Initialising...", flush=True)
|
||||
|
||||
self.de_client = DocumentEmbeddingsClient(
|
||||
pulsar_host=pulsar_host,
|
||||
subscriber=module + "-de",
|
||||
input_queue=de_request_queue,
|
||||
output_queue=de_response_queue,
|
||||
pulsar_api_key=pulsar_api_key,
|
||||
)
|
||||
|
||||
self.embeddings = EmbeddingsClient(
|
||||
pulsar_host=pulsar_host,
|
||||
input_queue=emb_request_queue,
|
||||
output_queue=emb_response_queue,
|
||||
subscriber=module + "-emb",
|
||||
pulsar_api_key=pulsar_api_key,
|
||||
)
|
||||
|
||||
self.lang = PromptClient(
|
||||
pulsar_host=pulsar_host,
|
||||
input_queue=pr_request_queue,
|
||||
output_queue=pr_response_queue,
|
||||
subscriber=module + "-de-prompt",
|
||||
pulsar_api_key=pulsar_api_key,
|
||||
)
|
||||
self.prompt_client = prompt_client
|
||||
self.embeddings_client = embeddings_client
|
||||
self.doc_embeddings_client = doc_embeddings_client
|
||||
|
||||
if self.verbose:
|
||||
print("Initialised", flush=True)
|
||||
|
||||
def query(
|
||||
async def query(
|
||||
self, query, user="trustgraph", collection="default",
|
||||
doc_limit=20,
|
||||
):
|
||||
|
|
@ -137,14 +75,17 @@ class DocumentRag:
|
|||
doc_limit=doc_limit
|
||||
)
|
||||
|
||||
docs = q.get_docs(query)
|
||||
docs = await q.get_docs(query)
|
||||
|
||||
if self.verbose:
|
||||
print("Invoke LLM...", flush=True)
|
||||
print(docs)
|
||||
print(query)
|
||||
|
||||
resp = self.lang.request_document_prompt(query, docs)
|
||||
resp = await self.prompt_client.document_prompt(
|
||||
query = query,
|
||||
documents = docs
|
||||
)
|
||||
|
||||
if self.verbose:
|
||||
print("Done", flush=True)
|
||||
|
|
|
|||
|
|
@ -5,88 +5,77 @@ Input is query, output is response.
|
|||
"""
|
||||
|
||||
from ... schema import DocumentRagQuery, DocumentRagResponse, Error
|
||||
from ... schema import document_rag_request_queue, document_rag_response_queue
|
||||
from ... schema import prompt_request_queue
|
||||
from ... schema import prompt_response_queue
|
||||
from ... schema import embeddings_request_queue
|
||||
from ... schema import embeddings_response_queue
|
||||
from ... schema import document_embeddings_request_queue
|
||||
from ... schema import document_embeddings_response_queue
|
||||
from ... log_level import LogLevel
|
||||
from ... document_rag import DocumentRag
|
||||
from ... base import ConsumerProducer
|
||||
from . document_rag import DocumentRag
|
||||
from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
|
||||
from ... base import PromptClientSpec, EmbeddingsClientSpec
|
||||
from ... base import DocumentEmbeddingsClientSpec
|
||||
|
||||
module = "document-rag"
|
||||
default_ident = "document-rag"
|
||||
|
||||
default_input_queue = document_rag_request_queue
|
||||
default_output_queue = document_rag_response_queue
|
||||
default_subscriber = module
|
||||
|
||||
class Processor(ConsumerProducer):
|
||||
class Processor(FlowProcessor):
|
||||
|
||||
def __init__(self, **params):
|
||||
|
||||
input_queue = params.get("input_queue", default_input_queue)
|
||||
output_queue = params.get("output_queue", default_output_queue)
|
||||
subscriber = params.get("subscriber", default_subscriber)
|
||||
pr_request_queue = params.get(
|
||||
"prompt_request_queue", prompt_request_queue
|
||||
)
|
||||
pr_response_queue = params.get(
|
||||
"prompt_response_queue", prompt_response_queue
|
||||
)
|
||||
emb_request_queue = params.get(
|
||||
"embeddings_request_queue", embeddings_request_queue
|
||||
)
|
||||
emb_response_queue = params.get(
|
||||
"embeddings_response_queue", embeddings_response_queue
|
||||
)
|
||||
de_request_queue = params.get(
|
||||
"document_embeddings_request_queue",
|
||||
document_embeddings_request_queue
|
||||
)
|
||||
de_response_queue = params.get(
|
||||
"document_embeddings_response_queue",
|
||||
document_embeddings_response_queue
|
||||
)
|
||||
id = params.get("id", default_ident)
|
||||
|
||||
doc_limit = params.get("doc_limit", 10)
|
||||
doc_limit = params.get("doc_limit", 5)
|
||||
|
||||
super(Processor, self).__init__(
|
||||
**params | {
|
||||
"input_queue": input_queue,
|
||||
"output_queue": output_queue,
|
||||
"subscriber": subscriber,
|
||||
"input_schema": DocumentRagQuery,
|
||||
"output_schema": DocumentRagResponse,
|
||||
"prompt_request_queue": pr_request_queue,
|
||||
"prompt_response_queue": pr_response_queue,
|
||||
"embeddings_request_queue": emb_request_queue,
|
||||
"embeddings_response_queue": emb_response_queue,
|
||||
"document_embeddings_request_queue": de_request_queue,
|
||||
"document_embeddings_response_queue": de_response_queue,
|
||||
"id": id,
|
||||
"doc_limit": doc_limit,
|
||||
}
|
||||
)
|
||||
|
||||
self.rag = DocumentRag(
|
||||
pulsar_host=self.pulsar_host,
|
||||
pulsar_api_key=self.pulsar_api_key,
|
||||
pr_request_queue=pr_request_queue,
|
||||
pr_response_queue=pr_response_queue,
|
||||
emb_request_queue=emb_request_queue,
|
||||
emb_response_queue=emb_response_queue,
|
||||
de_request_queue=de_request_queue,
|
||||
de_response_queue=de_response_queue,
|
||||
verbose=True,
|
||||
module=module,
|
||||
)
|
||||
|
||||
self.doc_limit = doc_limit
|
||||
|
||||
async def handle(self, msg):
|
||||
self.register_specification(
|
||||
ConsumerSpec(
|
||||
name = "request",
|
||||
schema = DocumentRagQuery,
|
||||
handler = self.on_request,
|
||||
)
|
||||
)
|
||||
|
||||
self.register_specification(
|
||||
EmbeddingsClientSpec(
|
||||
request_name = "embeddings-request",
|
||||
response_name = "embeddings-response",
|
||||
)
|
||||
)
|
||||
|
||||
self.register_specification(
|
||||
DocumentEmbeddingsClientSpec(
|
||||
request_name = "document-embeddings-request",
|
||||
response_name = "document-embeddings-response",
|
||||
)
|
||||
)
|
||||
|
||||
self.register_specification(
|
||||
PromptClientSpec(
|
||||
request_name = "prompt-request",
|
||||
response_name = "prompt-response",
|
||||
)
|
||||
)
|
||||
|
||||
self.register_specification(
|
||||
ProducerSpec(
|
||||
name = "response",
|
||||
schema = DocumentRagResponse,
|
||||
)
|
||||
)
|
||||
|
||||
async def on_request(self, msg, consumer, flow):
|
||||
|
||||
try:
|
||||
|
||||
self.rag = DocumentRag(
|
||||
embeddings_client = flow("embeddings-request"),
|
||||
doc_embeddings_client = flow("document-embeddings-request"),
|
||||
prompt_client = flow("prompt-request"),
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
v = msg.value()
|
||||
|
||||
# Sender-produced ID
|
||||
|
|
@ -99,11 +88,15 @@ class Processor(ConsumerProducer):
|
|||
else:
|
||||
doc_limit = self.doc_limit
|
||||
|
||||
response = self.rag.query(v.query, doc_limit=doc_limit)
|
||||
response = await self.rag.query(v.query, doc_limit=doc_limit)
|
||||
|
||||
print("Send response...", flush=True)
|
||||
r = DocumentRagResponse(response = response, error=None)
|
||||
await self.send(r, properties={"id": id})
|
||||
await flow("response").send(
|
||||
DocumentRagResponse(
|
||||
response = response,
|
||||
error = None
|
||||
),
|
||||
properties = {"id": id}
|
||||
)
|
||||
|
||||
print("Done.", flush=True)
|
||||
|
||||
|
|
@ -113,25 +106,21 @@ class Processor(ConsumerProducer):
|
|||
|
||||
print("Send error response...", flush=True)
|
||||
|
||||
r = DocumentRagResponse(
|
||||
error=Error(
|
||||
type = "llm-error",
|
||||
message = str(e),
|
||||
await flow("response").send(
|
||||
DocumentRagResponse(
|
||||
response = None,
|
||||
error = Error(
|
||||
type = "document-rag-error",
|
||||
message = str(e),
|
||||
),
|
||||
),
|
||||
response=None,
|
||||
properties = {"id": id}
|
||||
)
|
||||
|
||||
await self.send(r, properties={"id": id})
|
||||
|
||||
self.consumer.acknowledge(msg)
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
ConsumerProducer.add_args(
|
||||
parser, default_input_queue, default_subscriber,
|
||||
default_output_queue,
|
||||
)
|
||||
FlowProcessor.add_args(parser)
|
||||
|
||||
parser.add_argument(
|
||||
'-d', '--doc-limit',
|
||||
|
|
@ -140,43 +129,7 @@ class Processor(ConsumerProducer):
|
|||
help=f'Default document fetch limit (default: 10)'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--prompt-request-queue',
|
||||
default=prompt_request_queue,
|
||||
help=f'Prompt request queue (default: {prompt_request_queue})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--prompt-response-queue',
|
||||
default=prompt_response_queue,
|
||||
help=f'Prompt response queue (default: {prompt_response_queue})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--embeddings-request-queue',
|
||||
default=embeddings_request_queue,
|
||||
help=f'Embeddings request queue (default: {embeddings_request_queue})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--embeddings-response-queue',
|
||||
default=embeddings_response_queue,
|
||||
help=f'Embeddings response queue (default: {embeddings_response_queue})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--document-embeddings-request-queue',
|
||||
default=document_embeddings_request_queue,
|
||||
help=f'Document embeddings request queue (default: {document_embeddings_request_queue})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--document-embeddings-response-queue',
|
||||
default=document_embeddings_response_queue,
|
||||
help=f'Document embeddings response queue (default: {document_embeddings_response_queue})',
|
||||
)
|
||||
|
||||
def run():
|
||||
|
||||
Processor.launch(module, __doc__)
|
||||
Processor.launch(default_ident, __doc__)
|
||||
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@
|
|||
import asyncio
|
||||
|
||||
LABEL="http://www.w3.org/2000/01/rdf-schema#label"
|
||||
DEFINITION="http://www.w3.org/2004/02/skos/core#definition"
|
||||
|
||||
class Query:
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue