diff options
| -rw-r--r-- | rag/cli.py | 6 | ||||
| -rw-r--r-- | rag/generator/__init__.py | 2 | ||||
| -rw-r--r-- | rag/generator/abstract.py | 1 | ||||
| -rw-r--r-- | rag/generator/cohere.py | 9 | ||||
| -rw-r--r-- | rag/generator/ollama.py | 6 | ||||
| -rw-r--r-- | rag/generator/prompt.py | 1 | ||||
| -rw-r--r-- | rag/retriever/document.py | 2 | ||||
| -rw-r--r-- | rag/retriever/encoder.py | 2 | ||||
| -rw-r--r-- | rag/retriever/parser/pdf.py | 4 | ||||
| -rw-r--r-- | rag/retriever/retriever.py | 7 | ||||
| -rw-r--r-- | rag/ui.py | 6 |
11 files changed, 21 insertions, 25 deletions
@@ -1,10 +1,10 @@ from pathlib import Path + from dotenv import load_dotenv -from rag.generator import get_generator, MODELS -from rag.retriever.retriever import Retriever +from rag.generator import MODELS, get_generator from rag.generator.prompt import Prompt - +from rag.retriever.retriever import Retriever if __name__ == "__main__": load_dotenv() diff --git a/rag/generator/__init__.py b/rag/generator/__init__.py index 541eff8..ba23ffc 100644 --- a/rag/generator/__init__.py +++ b/rag/generator/__init__.py @@ -1,8 +1,8 @@ from typing import Type from .abstract import AbstractGenerator -from .ollama import Ollama from .cohere import Cohere +from .ollama import Ollama MODELS = ["ollama", "cohere"] diff --git a/rag/generator/abstract.py b/rag/generator/abstract.py index 5b336ea..1beacfb 100644 --- a/rag/generator/abstract.py +++ b/rag/generator/abstract.py @@ -1,5 +1,4 @@ from abc import abstractmethod - from typing import Any, Generator from .prompt import Prompt diff --git a/rag/generator/cohere.py b/rag/generator/cohere.py index cf95c18..7028b21 100644 --- a/rag/generator/cohere.py +++ b/rag/generator/cohere.py @@ -1,13 +1,12 @@ import os +from dataclasses import asdict from typing import Any, Generator -import cohere -from dataclasses import asdict +import cohere +from loguru import logger as log -from .prompt import Prompt from .abstract import AbstractGenerator - -from loguru import logger as log +from .prompt import Prompt class Cohere(metaclass=AbstractGenerator): diff --git a/rag/generator/ollama.py b/rag/generator/ollama.py index 1e968f6..cd8533b 100644 --- a/rag/generator/ollama.py +++ b/rag/generator/ollama.py @@ -4,11 +4,11 @@ from typing import Any, Generator, List import ollama from loguru import logger as log -from .prompt import Prompt -from .abstract import AbstractGenerator - from rag.retriever.vector import Document +from .abstract import AbstractGenerator +from .prompt import Prompt + SYSTEM_PROMPT = ( "# System Preamble" "## Basic Rules" diff --git a/rag/generator/prompt.py b/rag/generator/prompt.py index 059592a..41f264e 100644 --- a/rag/generator/prompt.py +++ b/rag/generator/prompt.py @@ -1,7 +1,6 @@ from dataclasses import dataclass from typing import List - from rag.retriever.vector import Document diff --git a/rag/retriever/document.py b/rag/retriever/document.py index 54ac451..67ff936 100644 --- a/rag/retriever/document.py +++ b/rag/retriever/document.py @@ -1,8 +1,8 @@ import hashlib import os -from langchain_community.document_loaders.blob_loaders import Blob import psycopg +from langchain_community.document_loaders.blob_loaders import Blob from loguru import logger as log TABLES = """ diff --git a/rag/retriever/encoder.py b/rag/retriever/encoder.py index 753157f..db69ee5 100644 --- a/rag/retriever/encoder.py +++ b/rag/retriever/encoder.py @@ -1,6 +1,6 @@ import os from pathlib import Path -from typing import List, Dict +from typing import Dict, List from uuid import uuid4 import ollama diff --git a/rag/retriever/parser/pdf.py b/rag/retriever/parser/pdf.py index 410f027..4c5addc 100644 --- a/rag/retriever/parser/pdf.py +++ b/rag/retriever/parser/pdf.py @@ -3,11 +3,11 @@ from pathlib import Path from typing import List, Optional from langchain.text_splitter import RecursiveCharacterTextSplitter -from langchain_core.documents import Document +from langchain_community.document_loaders.blob_loaders import Blob from langchain_community.document_loaders.parsers.pdf import ( PyPDFParser, ) -from langchain_community.document_loaders.blob_loaders import Blob +from langchain_core.documents import Document class PDFParser: diff --git a/rag/retriever/retriever.py b/rag/retriever/retriever.py index 885dafe..deffae5 100644 --- a/rag/retriever/retriever.py +++ b/rag/retriever/retriever.py @@ -1,12 +1,13 @@ +from io import BytesIO from pathlib import Path -from typing import Optional, List +from typing import List, Optional + from loguru import logger as log -from io import BytesIO from .document import DocumentDB from .encoder import Encoder from .parser.pdf import PDFParser -from .vector import VectorDB, Document +from .vector import Document, VectorDB class Retriever: @@ -1,13 +1,11 @@ import streamlit as st - +from dotenv import load_dotenv from langchain_community.document_loaders.blob_loaders import Blob -from dotenv import load_dotenv -from rag.generator import get_generator, MODELS +from rag.generator import MODELS, get_generator from rag.generator.prompt import Prompt from rag.retriever.retriever import Retriever - if __name__ == "__main__": load_dotenv() retriever = Retriever() |