llm - openai, gpt4all, llama

Author: apocas

.gitignore

@@ -162,4 +162,7 @@ cython_debug/
 embeddings/*
 !embeddings/.gitkeep
 projects/*
-!projects/.gitkeep
+!projects/.gitkeep
+models/*
+!models/.gitkeep
+.DS_Store

app/brain.py

@@ -1,13 +1,39 @@
 from fastapi import HTTPException
 from langchain.text_splitter import CharacterTextSplitter
+from langchain.chains import RetrievalQA
+from langchain import OpenAI
+from langchain.llms import GPT4All, LlamaCpp
 
 from app.project import Project
 
+LLMS = {
+    "openai": (OpenAI, {"temperature": 0, "model_name": "text-davinci-003"}),
+    "llamacpp": (LlamaCpp, {"model": "./models/ggml-model-q4_0.bin"}),
+    "gpt4all": (GPT4All, {"model": "./models/ggml-gpt4all-j-v1.3-groovy.bin", "backend": "gptj", "n_ctx": 1000}),
+}
+
+
 class Brain:
     def __init__(self):
-       self.projects = []
+        self.projects = []
+
+        self.text_splitter = CharacterTextSplitter(
+            separator=" ", chunk_size=1024, chunk_overlap=0)
+
+        self.llmCache = {}
 
-       self.text_splitter = CharacterTextSplitter(separator=" ", chunk_size=1024, chunk_overlap=0)
+    def loadLLM(self, llmModel, **kwargs):
+        if llmModel in self.llmCache:
+            return self.llmCache[llmModel]
+        else:
+            if llmModel in LLMS:
+                loader_class, llm_args = LLMS[llmModel]
+                llm = loader_class(**llm_args, **kwargs)
+                self.llmCache[llmModel] = llm
+                return llm
+            else:
+                raise HTTPException(
+                    status_code=500, detail='{"error": "Invalid LLM type."}')
 
     def listProjects(self):
         return [project.model.name for project in self.projects]
@@ -20,18 +46,31 @@ def createProject(self, projectModel):
 
     def loadProject(self, name):
         for project in self.projects:
-          if project.model.name == name:
-            return project
+            if project.model.name == name:
+                return project
 
         project = Project()
         project.load(name)
         self.projects.append(project)
         return project
-        
-    
+
     def deleteProject(self, name):
         for project in self.projects:
-          if project.model.name == name:
-            project.delete()
-            self.projects.remove(project)
+            if project.model.name == name:
+                project.delete()
+                self.projects.remove(project)
+
+    def question(self, project, questionModel):
+        retriever = project.db.as_retriever(
+            search_type="similarity", search_kwargs={"k": 2}
+        )
+
+        llm = self.loadLLM(questionModel.llm or project.model.llm)
+
+        qa = RetrievalQA.from_chain_type(
+            llm=llm,
+            chain_type="stuff",
+            retriever=retriever,
+        )
 
+        return qa.run(questionModel.question)

app/main.py

@@ -2,15 +2,14 @@
 import os
 from tempfile import NamedTemporaryFile
 from fastapi import FastAPI, HTTPException, Request, UploadFile
-from langchain import OpenAI
 from langchain.document_loaders import (
     WebBaseLoader,
 )
 from langchain.chains import RetrievalQA
 from dotenv import load_dotenv
 from app.brain import Brain
 
-from app.project import IngestModel, ProjectModel, QueryModel
+from app.models import IngestModel, ProjectModel, QuestionModel, ChatModel
 from app.tools import FindFileLoader, IndexDocuments
 
 load_dotenv()
@@ -118,26 +117,24 @@ def ingestFile(projectName: str, file: UploadFile):
         raise HTTPException(
             status_code=500, detail='{"error": ' + str(e) + '}')
 
-@app.post("/projects/{projectName}/query")
-def queryProject(projectName: str, input: QueryModel):
-    try:
-      project = brain.loadProject(projectName)
-
-      retriever = project.db.as_retriever(
-          search_type="similarity", search_kwargs={"k": 2}
-      )
-
-      llm = OpenAI(temperature=0, model_name="text-davinci-003") # type: ignore
 
-      qa = RetrievalQA.from_chain_type(
-          llm=llm,
-          chain_type="stuff",
-          retriever=retriever,
-      )
+@app.post("/projects/{projectName}/question")
+def questionProject(projectName: str, input: QuestionModel):
+    try:
+        project = brain.loadProject(projectName)
+        answer = brain.question(project, input)
+        return {"question": input.question, "answer": answer.strip()}
+    except Exception as e:
+        logging.error(e)
+        raise HTTPException(
+            status_code=500, detail='{"error": ' + str(e) + '}')
 
-      answer = qa.run(input.query)
 
-      return {"query": input.query, "answer": answer.strip()}
+@app.post("/projects/{projectName}/chat")
+def chatProject(projectName: str, input: ChatModel):
+    try:
+        project = brain.loadProject(projectName)
+        return "Not implemented yet."
     except Exception as e:
         raise HTTPException(
-            status_code=500, detail='{"error": ' + str(e) + '}')
+            status_code=500, detail='{"error": ' + str(e) + '}')

app/models.py

@@ -0,0 +1,21 @@
+from pydantic import BaseModel
+from typing import Union
+
+class IngestModel(BaseModel):
+    url: str
+
+class QuestionModel(BaseModel):
+    question: str
+    llm: Union[str, None] = None
+    
+class ChatModel(BaseModel):
+    message: str
+    conversation: Union[str, None] = None
+    llm: Union[str, None] = None
+
+class ProjectModel(BaseModel):
+    name: str
+    embeddings: Union[str, None] = None
+    embeddings_model: Union[str, None] = None
+    llm: Union[str, None] = None
+    llm_model: Union[str, None] = None

app/project.py

@@ -1,27 +1,13 @@
-from pydantic import BaseModel
-from typing import Union
 import json
 import os
 import shutil
 
 from langchain.vectorstores import Chroma
+from app.models import ProjectModel
 
 from app.tools import GetEmbedding
 
 
-class IngestModel(BaseModel):
-    url: str
-
-class QueryModel(BaseModel):
-    query: str
-
-class ProjectModel(BaseModel):
-    name: str
-    embeddings: Union[str, None] = None
-    embeddings_model: Union[str, None] = None
-    llm_model: Union[str, None] = None
-
-
 class Project:
 
     def boot(self, model: ProjectModel):

app/tools.py

@@ -24,7 +24,6 @@ def GetEmbedding(name: str, model=None):
 
     return embeddings[name]
 
-
 def IndexDocuments(brain, project, documents):
     texts = brain.text_splitter.split_documents(documents)
     texts_final = [doc.page_content for doc in texts]
@@ -40,7 +39,7 @@ def IndexDocuments(brain, project, documents):
 
 
 def FindFileLoader(temp, ext):
-    LOADERS_MAP = {
+    loaders = {
         ".csv": (CSVLoader, {}),
         ".doc": (UnstructuredWordDocumentLoader, {}),
         ".docx": (UnstructuredWordDocumentLoader, {}),
@@ -56,8 +55,8 @@ def FindFileLoader(temp, ext):
         ".txt": (TextLoader, {"encoding": "utf8"}),
     }
 
-    if ext in LOADERS_MAP:
-      loader_class, loader_args = LOADERS_MAP[ext]
+    if ext in loaders:
+      loader_class, loader_args = loaders[ext]
       return loader_class(temp.name, **loader_args)
     else:
         raise HTTPException(status_code=500, detail='{"error": "Invalid file type."}')

models/.gitkeep

@@ -6,4 +6,6 @@ python-multipart
 dotenv
 pdfminer.six
 unstructured
-pydantic
+pydantic
+pygpt4all
+llama-cpp-python

requirements.txt

@@ -0,0 +1,68 @@
+from fastapi.testclient import TestClient
+
+from app.main import app
+
+client = TestClient(app)
+
+
+def test_get():
+    response = client.get("/")
+    assert response.status_code == 200
+    assert response.json() == "REST AI API, so many 'A's and 'I's, so little time..."
+
+
+def test_createProjectHF():
+    response = client.post(
+        "/projects", json={"name": "test_huggingface",  "embeddings": "huggingface"})
+    assert response.status_code == 200
+
+
+def test_getProjectHF():
+    response = client.get("/projects/test_huggingface")
+    assert response.status_code == 200
+
+
+def test_ingestURLHF():
+    response = client.post("/projects/test_huggingface/ingest/url",
+                           json={"url": "https://www.google.com"})
+    assert response.status_code == 200
+
+
+def test_getProjectAfterIngestURLHF():
+    response = client.get("/projects/test_huggingface")
+    assert response.status_code == 200
+    assert response.json() == {
+        "project": "test_huggingface", "embeddings": "huggingface", "documents": 1, "metadatas": 1}
+
+
+def test_ingestUploadHF():
+    response = client.post("/projects/test_huggingface/ingest/upload",
+                           files={"file": ("test.txt", open("tests/test.txt", "rb"))})
+    assert response.status_code == 200
+
+
+def test_getProjectAfterIngestUploadHF():
+    response = client.get("/projects/test_huggingface")
+    assert response.status_code == 200
+    assert response.json() == {
+        "project": "test_huggingface", "embeddings": "huggingface", "documents": 2, "metadatas": 2}
+
+
+def test_questionProjectHF():
+    response = client.post("/projects/test_huggingface/question",
+                           json={"question": "What is the secret?", "llm": "gpt4all"})
+    assert response.status_code == 200
+    assert response.json() == {"question": "What is the secret?",
+                               "answer": "The secret is that ingenuity should be bigger than politics and corporate greed."}
+
+
+def test_deleteProjectHF():
+    response = client.delete("/projects/test_huggingface")
+    assert response.status_code == 200
+    assert response.json() == {"project": "test_huggingface"}
+
+
+def test_getProjectsAfterDelete():
+    response = client.get("/projects")
+    assert response.status_code == 200
+    assert response.json() == {"projects": []}

tests/test_huggingface.py

@@ -19,59 +19,56 @@ def test_getProjects():
 
 def test_createProject():
     response = client.post(
-        "/projects", json={"name": "test",  "embeddings": "openai"})
+        "/projects", json={"name": "test_openai",  "embeddings": "openai", "llm": "openai"})
     assert response.status_code == 200
 
-
 def test_getProject():
-    response = client.get("/projects/test")
+    response = client.get("/projects/test_openai")
     assert response.status_code == 200
 
 
 def test_ingestURL():
-    response = client.post("/projects/test/ingest/url",
+    response = client.post("/projects/test_openai/ingest/url",
                            json={"url": "https://www.google.com"})
     assert response.status_code == 200
 
 
 def test_getProjectAfterIngestURL():
-    response = client.get("/projects/test")
+    response = client.get("/projects/test_openai")
     assert response.status_code == 200
     assert response.json() == {
-        "project": "test", "embeddings": "openai", "documents": 1, "metadatas": 1}
+        "project": "test_openai", "embeddings": "openai", "documents": 1, "metadatas": 1}
 
 
 def test_ingestUpload():
-    response = client.post("/projects/test/ingest/upload",
+    response = client.post("/projects/test_openai/ingest/upload",
                            files={"file": ("test.txt", open("tests/test.txt", "rb"))})
     assert response.status_code == 200
 
 
 def test_getProjectAfterIngestUpload():
-    response = client.get("/projects/test")
+    response = client.get("/projects/test_openai")
     assert response.status_code == 200
     assert response.json() == {
-        "project": "test", "embeddings": "openai", "documents": 2, "metadatas": 2}
+        "project": "test_openai", "embeddings": "openai", "documents": 2, "metadatas": 2}
 
 
-def test_query():
-    response = client.post("/projects/test/query",
-                           json={"query": "What is the secret?"})
+def test_questionProject():
+    response = client.post("/projects/test_openai/question",
+                           json={"question": "What is the secret?"})
     assert response.status_code == 200
-    assert response.json() == {"query": "What is the secret?",
+    assert response.json() == {"question": "What is the secret?",
                                "answer": "The secret is that ingenuity should be bigger than politics and corporate greed."}
 
-
 def test_deleteProject():
-    response = client.delete("/projects/test")
+    response = client.delete("/projects/test_openai")
     assert response.status_code == 200
-    assert response.json() == {"project": "test"}
-
+    assert response.json() == {"project": "test_openai"}
 
 def test_getProjectAfterDelete():
-    response = client.get("/projects/test")
+    response = client.get("/projects/test_openai")
     assert response.status_code == 404
-
+    
 
 def test_getProjectsAfterDelete():
     response = client.get("/projects")