diff options
Diffstat (limited to 'rag/llm')
| -rw-r--r-- | rag/llm/encoder.py | 11 | ||||
| -rw-r--r-- | rag/llm/generator.py | 7 |
2 files changed, 10 insertions, 8 deletions
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) |