From ab616b27791ba5d5ab112984ee55b8e9ee9fe729 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Wed, 17 Jul 2024 16:55:46 +0100
Subject: [PATCH] Fix bugs

---
 trustgraph/chunker/recursive/chunker.py        | 3 +++
 trustgraph/decoder/pdf/pdf_decoder.py          | 2 ++
 trustgraph/kg/extract_definitions/extract.py   | 3 +++
 trustgraph/kg/extract_relationships/extract.py | 4 ++--
 trustgraph/rag/graph/rag.py                    | 4 ++--
 5 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/trustgraph/chunker/recursive/chunker.py b/trustgraph/chunker/recursive/chunker.py
index b26ed70b..b116eca5 100755
--- a/trustgraph/chunker/recursive/chunker.py
+++ b/trustgraph/chunker/recursive/chunker.py
@@ -4,6 +4,9 @@ Simple decoder, accepts text documents on input, outputs chunks from the
 as text as separate output objects.
 """
 
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+
 from ... schema import TextDocument, Chunk, Source
 from ... log_level import LogLevel
 from ... base import ConsumerProducer
diff --git a/trustgraph/decoder/pdf/pdf_decoder.py b/trustgraph/decoder/pdf/pdf_decoder.py
index dd241589..e87f2905 100755
--- a/trustgraph/decoder/pdf/pdf_decoder.py
+++ b/trustgraph/decoder/pdf/pdf_decoder.py
@@ -4,6 +4,8 @@ Simple decoder, accepts PDF documents on input, outputs pages from the
 PDF document as text as separate output objects.
 """
 
+import tempfile
+import base64
 from langchain_community.document_loaders import PyPDFLoader
 
 from ... schema import Document, TextDocument, Source
diff --git a/trustgraph/kg/extract_definitions/extract.py b/trustgraph/kg/extract_definitions/extract.py
index c86e926a..42cbacc9 100755
--- a/trustgraph/kg/extract_definitions/extract.py
+++ b/trustgraph/kg/extract_definitions/extract.py
@@ -4,6 +4,9 @@ Simple decoder, accepts vector+text chunks input, applies entity analysis to
 get entity definitions which are output as graph edges.
 """
 
+import urllib.parse
+import json
+
 from ... schema import VectorsChunk, Triple, Source, Value
 from ... log_level import LogLevel
 from ... llm_client import LlmClient
diff --git a/trustgraph/kg/extract_relationships/extract.py b/trustgraph/kg/extract_relationships/extract.py
index e2aaeef1..adc0b71a 100755
--- a/trustgraph/kg/extract_relationships/extract.py
+++ b/trustgraph/kg/extract_relationships/extract.py
@@ -5,6 +5,8 @@ relationship analysis to get entity relationship edges which are output as
 graph edges.
 """
 
+import urllib.parse
+import json
 from pulsar.schema import JsonSchema
 
 from ... schema import VectorsChunk, Triple, VectorsAssociation, Source, Value
@@ -84,8 +86,6 @@ class Processor(ConsumerProducer):
 
         chunk = v.chunk.decode("utf-8")
 
-        g = rdflib.Graph()
-
         try:
 
             rels = self.get_relationships(chunk)
diff --git a/trustgraph/rag/graph/rag.py b/trustgraph/rag/graph/rag.py
index 4a9b820e..8a7a484d 100755
--- a/trustgraph/rag/graph/rag.py
+++ b/trustgraph/rag/graph/rag.py
@@ -37,8 +37,8 @@ class Processor(ConsumerProducer):
             input_queue=input_queue,
             output_queue=output_queue,
             subscriber=subscriber,
-            input_schema=TextCompletionRequest,
-            output_schema=TextCompletionResponse,
+            input_schema=GraphRagQuery,
+            output_schema=GraphRagResponse,
         )
 
         self.rag = GraphRag(