wip: postgres vecto.rs db, s/haystack/langchain

.gitignore

@@ -27,3 +27,6 @@ __pycache__/
 
 # macOS
 .DS_Store
+
+# logs
+*.log

Dockerfile

@@ -12,18 +12,31 @@ FROM $CUDA_FROM
 
 ENV PATH="/usr/local/cuda/bin:${PATH}"
 
+# Install unattendedly
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Force a config for tzdata package, otherwise it will interactively ask during install
+RUN ln -fs /usr/share/zoneinfo/UTC /etc/localtime
+
 # Install essential packages from ubuntu repository
 RUN apt-get update -y && \
     apt-get install -y --no-install-recommends openssh-server openssh-client git git-lfs && \
     apt-get install -y curl && \
     apt-get install -y python3 python3-pip python3-venv && \
+    apt-get install -y postgresql-14 && \
+    apt-get install -y jq && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
 
+# Install vecto.rs extension to postgres
+RUN curl -L -O https://github.com/tensorchord/pgvecto.rs/releases/download/v0.2.0/vectors-pg14_0.2.0_amd64.deb
+RUN dpkg -i vectors-pg14_0.2.0_amd64.deb
+
+
 # Install node from upstream, ubuntu packages are too old
-RUN curl -sL https://deb.nodesource.com/setup_18.x | bash
-RUN apt-get install -y nodejs && \
+RUN curl -sL https://deb.nodesource.com/setup_18.x | bash && \
+    apt-get install -y nodejs && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
@@ -53,11 +66,13 @@ ARG OLLAMA_URL=http://localhost:11434
 ENV OLLAMA_MODEL_NAME=${OLLAMA_MODEL_NAME}
 ENV OLLAMA_URL=${OLLAMA_URL}
 
+# TODO: cache path
 RUN ollama serve & while ! curl ${OLLAMA_URL}; do sleep 1; done; ollama pull $OLLAMA_MODEL_NAME
 
 
 # Load sentence-transformers model once in order to cache it in the image
 # TODO: ARG / ENV for embedder model
+# TODO: SENTENCE_TRANSFORMERS_HOME for cache path
 RUN echo "from haystack.components.embedders import SentenceTransformersDocumentEmbedder\nSentenceTransformersDocumentEmbedder(model='svalabs/german-gpl-adapted-covid').warm_up()" | python3
 
 

README.md

@@ -10,9 +10,7 @@ To build and run the container locally with hot reload on python files do:
 ```
 DOCKER_BUILDKIT=1 docker build . -t gbnc
 docker run  \
-  --env HUGGING_FACE_HUB_TOKEN=$HUGGING_FACE_HUB_TOKEN \
   --volume "$(pwd)/gswikichat":/workspace/gswikichat \
-  --volume gbnc_cache:/root/.cache \
   --publish 8000:8000 \
   --rm \
   --interactive \
@@ -22,6 +20,21 @@ docker run  \
 ```
 Point your browser to http://localhost:8000/ and use the frontend.
 
+To fetch data from a `toc.json` wiki fetching definition, run:
+```
+$ docker exec -it gbnc bash
+# export WIKI_USER=<wikibotusername>
+# export WIKI_PW=<yoursecretbotuserpassword>
+# python3 -m gswikichat.fetch_articles toc.json > articles.json
+```
+
+To import data run:
+```
+$ docker exec -it gbnc bash
+# cat json_input/excellent-articles_10.json | jq 'to_entries | map({content: .value, meta: {source: .key}})' > import.json
+# python3 -m gswikichat.db import.json
+```
+
 ### Runpod.io
 
 The container works on [runpod.io](https://www.runpod.io/) GPU instances. A [template is available here](https://runpod.io/gsc?template=0w8z55rf19&ref=yfvyfa0s).

frontend/src/components/field/FieldAnswer.vue

@@ -17,8 +17,8 @@
               class="text-sm cursor-pointer text-light-distinct-text dark:text-dark-distinct-text"
             >
               <summary>
-                {{ $t('source') }} ({{ s.score.toFixed(1) }}/5):
-                <a class="link-text" :href="s.src">{{ s.src }}</a>
+                {{ $t('source') }} ({{ s.score.toFixed(1) }}):
+                <a class="link-text" :href="s.source">{{ s.source }}</a>
               </summary>
               <p class="pt-2 pl-4">{{ s.content }}</p>
             </details>

frontend/src/types/source.d.ts

@@ -1,6 +1,6 @@
 export type Source = {
   id: number
-  src: string
+  source: string
   content: string
   score: number
 }

gswikichat/api.py

@@ -6,61 +6,57 @@
 from fastapi.staticfiles import StaticFiles
 from fastapi import FastAPI, Header
 
-from .rag import rag_pipeline
-
 from .logger import get_logger
+from .rag import rag_pipeline
 
 # Create logger instance from base logger config in `logger.py`
 logger = get_logger(__name__)
 
-FRONTEND_STATIC_DIR = './frontend/dist'
+FRONTEND_STATIC_DIR = "./frontend/dist"
 API_SECRET = os.environ.get("API_SECRET")
 
 app = FastAPI()
 
 app.mount(
     "/assets",
     StaticFiles(directory=f"{FRONTEND_STATIC_DIR}/assets"),
-    name="frontend-assets"
+    name="frontend-assets",
 )
 
+
 @app.get("/")
 async def root():
     return FileResponse(f"{FRONTEND_STATIC_DIR}/index.html")
 
+
 @app.get("/favicon.ico")
 async def favicon():
     return FileResponse(f"{FRONTEND_STATIC_DIR}/favicon.ico")
 
+
 @app.get("/api")
 async def api(x_api_secret: Annotated[str, Header()], query, top_k=3, lang='en'):
     if not API_SECRET == x_api_secret:
         raise Exception("API key is missing or incorrect") 
 
-    if not lang in ['en', 'de']:
+    if not lang in ["en", "de"]:
         raise Exception("language must be 'en' or 'de'")
 
-    logger.debug(f'{query=}')  # Assuming we change the input name
-    logger.debug(f'{top_k=}')
-    logger.debug(f'{lang=}')
+    logger.debug(f"{query=}")
+    logger.debug(f"{top_k=}")
+    logger.debug(f"{lang=}")
+
+    answer = rag_pipeline(query=query, top_k=top_k, lang=lang)
 
-    answer = rag_pipeline(
-        query=query,
-        top_k=top_k,
-        lang=lang
-    )
+    if not answer:
+        return {}
 
     sources = [
-        {
-            "src": d_.meta['src'],
-            "content": d_.content,
-            "score": d_.score
-        } for d_ in answer.documents
+        {"id": d_.id, "source": d_.meta["source"], "content": d_.content, "score": d_.score}
+        for d_ in answer.documents
     ]
 
-    logger.debug(f'{answer=}')
+    logger.debug(f"{answer.data=}")
+    logger.debug(f"{answer.documents=}")
 
-    return {
-        "answer": answer.data.content,
-        "sources": sources
-    }
+    return {"answer": answer.data.content, "sources": sources}

gswikichat/db.py

@@ -0,0 +1,107 @@
+import os
+
+import torch
+
+from langchain.text_splitter import CharacterTextSplitter
+from langchain_community.document_loaders import JSONLoader
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores.pgvecto_rs import PGVecto_rs
+
+from .logger import get_logger
+
+
+SENTENCE_TRANSFORMER_MODEL = "svalabs/german-gpl-adapted-covid"
+
+logger = get_logger(__name__)
+
+
+def get_device():
+    device = "cpu"
+    if torch.cuda.is_available():
+        logger.info("GPU is available.")
+        device = "cuda"
+    return device
+
+
+def get_embedding_model():
+    # https://huggingface.co/svalabs/german-gpl-adapted-covid
+    logger.info(f"Embedding model: {SENTENCE_TRANSFORMER_MODEL}")
+
+    return HuggingFaceEmbeddings(
+        model_name=SENTENCE_TRANSFORMER_MODEL,
+        model_kwargs={"device": get_device()},
+        show_progress=True,
+    )
+
+
+def get_db():
+    PORT = os.getenv("DB_PORT", 5432)
+    HOST = os.getenv("DB_HOST", "127.0.0.1")
+    USER = os.getenv("DB_USER", "gbnc")
+    PASS = os.getenv("DB_PASS", "")
+    DB_NAME = os.getenv("DB_NAME", "gbnc")
+
+    URL = "postgresql+psycopg://{username}:{password}@{host}:{port}/{db_name}".format(
+        port=PORT,
+        host=HOST,
+        username=USER,
+        password=PASS,
+        db_name=DB_NAME,
+    )
+
+    return PGVecto_rs.from_collection_name(
+        embedding=get_embedding_model(),
+        db_url=URL,
+        collection_name="gbnc",
+    )
+
+
+def import_data(file):
+    def metadata_func(record: dict, metadata: dict) -> dict:
+        metadata["source"] = record.get("meta", {}).get("source")
+        return metadata
+
+    loader = JSONLoader(
+        file_path=file,
+        jq_schema=".[]",
+        content_key="content",
+        metadata_func=metadata_func,
+    )
+
+    documents = loader.load()
+
+    logger.debug(f"Loaded {len(documents)} documents.")
+
+    text_splitter = CharacterTextSplitter(chunk_size=250, chunk_overlap=0)
+    chunks = text_splitter.split_documents(documents)
+    logger.debug(f"Split documents into {len(chunks)} chunks.")
+
+    logger.debug(f"Importing into database.")
+    get_db().add_documents(chunks)
+
+
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) > 1:
+        file = sys.argv[1]
+        import_data(file)
+
+    else:
+        logger.error(
+            """Provide JSON file with the following structure as first parameter
+    [
+        {
+            "content":"document content one", "meta":{
+                "source": "https://source.url/one"
+            }
+        },
+        {
+            "content":"document content two", "meta":{
+                "source": "https://source.url/two"
+            }
+        }
+    ]
+            """
+        )
+        sys.exit(1)