From 8211705debf9d1335223c606275f46c43c78d8a2 Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Mon, 8 Apr 2024 00:23:52 +0200
Subject: Updates

---
 rag/llm/encoder.py   | 11 +++++++----
 rag/llm/generator.py |  7 +++----
 2 files changed, 10 insertions(+), 8 deletions(-)

(limited to 'rag/llm')

diff --git a/rag/llm/encoder.py b/rag/llm/encoder.py
index 7c03dc5..95f3c6a 100644
--- a/rag/llm/encoder.py
+++ b/rag/llm/encoder.py
@@ -1,5 +1,5 @@
 import os
-from typing import List
+from typing import Iterator, List
 from uuid import uuid4
 
 import ollama
@@ -7,8 +7,11 @@ from langchain_core.documents import Document
 from loguru import logger as log
 from qdrant_client.http.models import StrictFloat
 
-from rag.db.vector import Point
 
+try:
+    from rag.db.vector import Point
+except ModuleNotFoundError:
+    from db.vector import Point
 
 class Encoder:
     def __init__(self) -> None:
@@ -18,11 +21,11 @@ class Encoder:
     def __encode(self, prompt: str) -> List[StrictFloat]:
         return list(ollama.embeddings(model=self.model, prompt=prompt)["embedding"])
 
-    def encode_document(self, chunks: List[Document]) -> List[Point]:
+    def encode_document(self, chunks: Iterator[Document]) -> List[Point]:
         log.debug("Encoding document...")
         return [
             Point(
-                id=str(uuid4()),
+                id=uuid4().hex,
                 vector=self.__encode(chunk.page_content),
                 payload={"text": chunk.page_content},
             )
diff --git a/rag/llm/generator.py b/rag/llm/generator.py
index b0c6c40..8c7702f 100644
--- a/rag/llm/generator.py
+++ b/rag/llm/generator.py
@@ -15,9 +15,8 @@ class Generator:
     def __init__(self) -> None:
         self.model = os.environ["GENERATOR_MODEL"]
 
-    def __metaprompt(self, role: str, prompt: Prompt) -> str:
+    def __metaprompt(self, prompt: Prompt) -> str:
         metaprompt = (
-            f"You are a {role}.\n"
             "Answer the following question using the provided context.\n"
             "If you can't find the answer, do not pretend you know it,"
             'but answer "I don\'t know".\n\n'
@@ -28,7 +27,7 @@ class Generator:
         )
         return metaprompt
 
-    def generate(self, prompt: Prompt, role: str) -> str:
+    def generate(self, prompt: Prompt) -> str:
         log.debug("Generating answer...")
-        metaprompt = self.__metaprompt(role, prompt)
+        metaprompt = self.__metaprompt(prompt)
         return ollama.generate(model=self.model, prompt=metaprompt)
-- 
cgit v1.2.3-70-g09d2