Cool poc

.devcontainer/devcontainer.json

@@ -0,0 +1,44 @@
+// For format details, see https://aka.ms/devcontainer.json. For config options, see the
+// README at: https://github.com/devcontainers/templates/tree/main/src/docker-existing-docker-compose
+{
+	"name": "Existing Docker Compose (Extend)",
+
+	// Update the 'dockerComposeFile' list if you have more compose files or use different names.
+	// The .devcontainer/docker-compose.yml file contains any overrides you need/want to make.
+	"dockerComposeFile": [
+		"../docker-compose.yml",
+		"docker-compose.yml"
+	],
+
+	// The 'service' property is the name of the service for the container that VS Code should
+	// use. Update this value and .devcontainer/docker-compose.yml to the real service name.
+	"service": "python-loader",
+
+	// The optional 'workspaceFolder' property is the path VS Code should open by default when
+	// connected. This is typically a file mount in .devcontainer/docker-compose.yml
+	"workspaceFolder": "/workspaces/${localWorkspaceFolderBasename}",
+	"features": {
+		"ghcr.io/devcontainers/features/python:1": {}
+	},
+
+	// Features to add to the dev container. More info: https://containers.dev/features.
+	// "features": {},
+
+	// Use 'forwardPorts' to make a list of ports inside the container available locally.
+	// "forwardPorts": [],
+
+	// Uncomment the next line if you want start specific services in your Docker Compose config.
+	// "runServices": [],
+
+	// Uncomment the next line if you want to keep your containers running after VS Code shuts down.
+	// "shutdownAction": "none",
+
+	// Uncomment the next line to run commands after the container is created.
+	// "postCreateCommand": "cat /etc/os-release",
+
+	// Configure tool-specific properties.
+	// "customizations": {},
+
+	// Uncomment to connect as an existing user other than the container default. More info: https://aka.ms/dev-containers-non-root.
+	"remoteUser": "root"
+}

.devcontainer/docker-compose.yml

@@ -0,0 +1,26 @@
+version: '3.8'
+services:
+  # Update this to the name of the service you want to work with in your docker-compose.yml file
+  python-loader:
+    # Uncomment if you want to override the service's Dockerfile to one in the .devcontainer 
+    # folder. Note that the path of the Dockerfile and context is relative to the *primary* 
+    # docker-compose.yml file (the first in the devcontainer.json "dockerComposeFile"
+    # array). The sample below assumes your primary file is in the root of your project.
+    #
+    # build:
+    #   context: .
+    #   dockerfile: .devcontainer/Dockerfile
+
+    volumes:
+      # Update this to wherever you want VS Code to mount the folder of your project
+      - ..:/workspaces:cached
+
+    # Uncomment the next four lines if you will use a ptrace-based debugger like C++, Go, and Rust.
+    # cap_add:
+    #   - SYS_PTRACE
+    # security_opt:
+    #   - seccomp:unconfined
+
+    # Overrides default command so things don't shut down after the process ends.
+    command: sleep infinity
+ 

.gitignore

@@ -0,0 +1,2 @@
+ollama_models
+chroma_db_data

README.md

@@ -3,8 +3,28 @@
 
 ## Getting started
 
-Start a developer console in NixOS with all need dependencies
+Download some models
 
 ```bash
-nix develop
+ollama pull mxbai-embed-large # Used for embeddings
+ollama pull gemma3 # Used as LLM
+```
+
+Download python depedencies
+
+```bash
+pip install -r requirements.txt
+```
+## Run
+
+First you need to seed the database with a few documents
+
+```bash
+python script.py seed
+```
+
+And then you can do your search
+
+```bash
+python script.py search --query "How does AI change industries?"
 ```

docker-compose.yml

@@ -1,8 +1,6 @@
 services:
   chroma:
     image: chromadb/chroma:latest
-    ports:
-      - "8000:8000"
     volumes:
       - ./chroma_db_data:/db
     environment:
@@ -11,6 +9,8 @@ services:
       - POSTGRES_USER=chroma
       - POSTGRES_PASSWORD=chroma
       - POSTGRES_DB=chroma
+      - CHROMA_SERVER_HOST=0.0.0.0
+      - CHROMA_SERVER_HTTP_PORT=8000
     depends_on:
       - chroma-db
 
@@ -31,3 +31,6 @@ services:
       - ./ollama_models:/root/.ollama
     environment:
       - OLLAMA_HOST=0.0.0.0
+
+  python-loader:
+    image: mcr.microsoft.com/devcontainers/base:jammy

flake.lock

@@ -1,26 +0,0 @@
-{
-  "nodes": {
-    "nixpkgs": {
-      "locked": {
-        "lastModified": 1743315132,
-        "narHash": "sha256-6hl6L/tRnwubHcA4pfUUtk542wn2Om+D4UnDhlDW9BE=",
-        "owner": "NixOS",
-        "repo": "nixpkgs",
-        "rev": "52faf482a3889b7619003c0daec593a1912fddc1",
-        "type": "github"
-      },
-      "original": {
-        "id": "nixpkgs",
-        "ref": "nixos-unstable",
-        "type": "indirect"
-      }
-    },
-    "root": {
-      "inputs": {
-        "nixpkgs": "nixpkgs"
-      }
-    }
-  },
-  "root": "root",
-  "version": 7
-}

flake.nix

@@ -1,21 +0,0 @@
-{
-  description = "A flake for Terraform development";
-
-  inputs.nixpkgs.url = "nixpkgs/nixos-unstable";
-
-  outputs = { self, nixpkgs }: {
-    devShells.x86_64-linux.default =
-      let
-        system = "x86_64-linux";
-        pkgs = import nixpkgs {
-          inherit system;
-          config.allowUnfree = true;
-        };
-      in
-      pkgs.mkShell {
-        buildInputs = with pkgs; [
-          docker-compose
-        ];
-      };
-  };
-}

requirements.txt

@@ -0,0 +1,2 @@
+chromadb
+requests

script.py

@@ -0,0 +1,86 @@
+import chromadb
+import requests
+import argparse
+
+OLLAMA_URL = "http://ollama:11434/api"
+CHROMA_COLLECTION_NAME = "rag_documents"
+
+# Connect to ChromaDB client
+client = chromadb.HttpClient(host="chroma", port=8000)
+collection = client.get_or_create_collection(name=CHROMA_COLLECTION_NAME)
+
+
+def get_embedding(text, embedding_model):
+    """Generate an embedding using Ollama."""
+    response = requests.post(f"{OLLAMA_URL}/embeddings", json={"model": embedding_model, "prompt": text})
+    if response.status_code == 200:
+        return response.json().get("embedding")
+    else:
+        raise Exception(f"Failed to get embedding: {response.text}")
+
+def seed_database(embedding_model):
+    """Seed the ChromaDB with example documents."""
+    documents = [
+        "Artificial Intelligence is transforming industries.",
+        "Machine learning helps computers learn from data.",
+        "Deep learning uses neural networks to process information.",
+        "RAG enhances language models with retrieved knowledge."
+    ]
+
+    for i, doc in enumerate(documents):
+        embedding = get_embedding(doc, embedding_model)
+        collection.add(ids=[str(i)], embeddings=[embedding], metadatas=[{"text": doc}])
+        print(f"Added document {i}: {doc}")
+
+    print("Database seeding complete.")
+
+def search(query, embedding_model, llm_model, top_k=2):
+    """Retrieve similar documents and generate an answer using an LLM."""
+    query_embedding = get_embedding(query, embedding_model)
+    results = collection.query(query_embeddings=[query_embedding], n_results=top_k)
+
+    retrieved_docs = [doc["text"] for doc_list in results["metadatas"] for doc in doc_list]
+    
+    if not retrieved_docs:
+        print("No relevant documents found.")
+        return
+
+    # Construct the LLM prompt
+    prompt = f"Use the following documents to answer the question:\n\n"
+    for doc in retrieved_docs:
+        prompt += f"- {doc}\n"
+    prompt += f"\nQuestion: {query}\nAnswer:"
+
+    response = requests.post(f"{OLLAMA_URL}/generate", json={"model": llm_model, "prompt": prompt, "stream": False})
+
+    print("RAW RESPONSE:", response.text)
+
+    if response.status_code == 200:
+        try:
+            data = response.json()
+            answer = data.get("response", "No response field found.")
+        except requests.exceptions.JSONDecodeError:
+            answer = response.text  # Fallback to raw text
+        print("\nSearch Results:\n")
+        for doc in retrieved_docs:
+            print(f"Retrieved: {doc}")
+        print("\nGenerated Answer:\n", answer)
+    else:
+        print(f"Failed to generate response: {response.status_code} - {response.text}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("command", choices=["seed", "search"], help="Command to run")
+    parser.add_argument("--query", type=str, help="Query text for searching")
+    parser.add_argument("--embedding_model", type=str, default="mxbai-embed-large", help="Embedding model")
+    parser.add_argument("--llm_model", type=str, default="gemma3", help="LLM model for generating responses")
+
+    args = parser.parse_args()
+
+    if args.command == "seed":
+        seed_database(args.embedding_model)
+    elif args.command == "search":
+        if not args.query:
+            print("Please provide a query with --query")
+        else:
+            search(args.query, args.embedding_model, args.llm_model)