RAG MVP

apps/experiments/views/experiment.py

@@ -24,7 +24,10 @@
 from django.views.decorators.http import require_POST
 from django.views.generic import CreateView, UpdateView
 from django_tables2 import SingleTableView
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import PyMuPDFLoader, TextLoader
 from langchain_core.prompts import PromptTemplate
+from langchain_openai import OpenAIEmbeddings
 from waffle import flag_is_active
 
 from apps.annotations.models import Tag
@@ -61,6 +64,7 @@
 from apps.teams.decorators import login_and_team_required
 from apps.teams.mixins import LoginAndTeamRequiredMixin
 from apps.users.models import CustomUser
+from apps.vectordb.vectorstore import PGVector
 
 
 @login_and_team_required
@@ -361,8 +365,40 @@ def form_valid(self, form):
         experiment = get_object_or_404(Experiment, team=self.request.team, pk=self.kwargs["pk"])
         file = super().form_valid(form)
         experiment.files.add(file)
+        file_path = experiment.files.all().last().file.path
+        splits = self.load_rag_file(file_path)
+        embeddings_model = OpenAIEmbeddings()
+        PGVector.from_texts(splits, embeddings_model, None, experiment)
         return file
 
+    def load_rag_file(self, file_path: str) -> list[str]:
+        """
+        Loads a text file of any supported type (PDF, TXT, HTML) into Langchain.
+
+        Args:
+            file_path (str): The path to the text file.
+
+        Returns:
+            str_splits: A list of strings from  Langchain Document objects
+            containing the loaded page_content.
+        """
+
+        # Automatically detect loader based on file extension if not provided
+        extension = file_path.split(".")[-1].lower()
+        if extension == "pdf":
+            loader = PyMuPDFLoader(file_path, extract_images=False)
+        elif extension in ("txt", "text"):
+            loader = TextLoader(file_path)
+        else:
+            raise ValueError(f"Unsupported file type: {extension}")
+
+        # Load the text file using the appropriate loader
+        documents = loader.load()
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        splits = text_splitter.split_documents(documents)
+        str_splits = [s.page_content for s in splits[0:10]]
+        return str_splits
+
     def get_delete_url(self, file):
         return reverse("experiments:remove_file", args=[self.request.team.slug, self.kwargs["pk"], file.pk])
 

apps/service_providers/llm_service/main.py

@@ -7,6 +7,7 @@
 from langchain_community.chat_models import ChatAnthropic
 from langchain_core.callbacks import BaseCallbackHandler
 from langchain_core.language_models import BaseLanguageModel
+from langchain_openai import OpenAIEmbeddings
 from langchain_openai.chat_models import AzureChatOpenAI, ChatOpenAI
 from openai import OpenAI
 from openai._base_client import SyncAPIClient
@@ -42,6 +43,14 @@ class OpenAILlmService(LlmService):
     openai_api_base: str = None
     openai_organization: str = None
 
+    def get_openai_embeddings(self, model="text-embedding-3-small") -> OpenAIEmbeddings:
+        return OpenAIEmbeddings(
+            openai_api_key=self.openai_api_key,
+            openai_api_base=self.openai_api_base,
+            openai_organization=self.openai_organization,
+            model=model,
+        )
+
     def get_raw_client(self) -> OpenAI:
         return OpenAI(api_key=self.openai_api_key, organization=self.openai_organization, base_url=self.openai_api_base)
 

apps/service_providers/llm_service/runnables.py

@@ -9,6 +9,7 @@
 
 import openai
 import pytz
+from langchain import hub
 from langchain.agents import AgentExecutor, create_openai_tools_agent
 from langchain.agents.openai_assistant.base import OpenAIAssistantFinish
 from langchain.memory import ConversationBufferMemory
@@ -29,6 +30,7 @@
 from apps.chat.conversation import compress_chat_history
 from apps.chat.models import ChatMessage, ChatMessageType
 from apps.experiments.models import Experiment, ExperimentSession
+from apps.vectordb.vectorstore import PGVector
 
 logger = logging.getLogger(__name__)
 
@@ -51,7 +53,8 @@ def create_experiment_runnable(experiment: Experiment, session: ExperimentSessio
     assert experiment.llm_provider, "Experiment must have an LLM provider"
     if experiment.tools_enabled:
         return AgentExperimentRunnable(experiment=experiment, session=session)
-
+    if experiment.files.all():
+        return RagExperimentRunnable(experiment=experiment, session=session)
     return SimpleExperimentRunnable(experiment=experiment, session=session)
 
 
@@ -219,6 +222,26 @@ def _build_chain(self) -> Runnable[dict[str, Any], str]:
         )
 
 
+class RagExperimentRunnable(ExperimentRunnable):
+    def _build_chain(self) -> Runnable[dict[str, Any], str]:
+        def format_docs(docs):
+            return "\n\n".join(doc.page_content for doc in docs)
+
+        model = self.llm_service.get_chat_model(self.experiment.llm, self.experiment.temperature)
+        embeddings = self.experiment.get_llm_service().get_openai_embeddings()
+        retriever = PGVector(self.experiment, embeddings).as_retriever()
+        prompt = hub.pull("rlm/rag-prompt")
+        return (
+            {"context": retriever | format_docs, "question": RunnablePassthrough()}
+            | RunnablePassthrough.assign(
+                history=RunnableLambda(self.memory.load_memory_variables) | itemgetter("history")
+            )
+            | prompt
+            | model
+            | StrOutputParser()
+        )
+
+
 class AgentExperimentRunnable(ExperimentRunnable):
     def _parse_output(self, output):
         return output.get("output", "")

requirements/requirements.in

@@ -45,3 +45,5 @@ turn-python>=0.2.0
 jinja2
 django-taggit
 pgvector
+PyMuPDF
+langchainhub