diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json new file mode 100644 index 0000000..5ac6e7d --- /dev/null +++ b/.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" +} diff --git a/.devcontainer/docker-compose.yml b/.devcontainer/docker-compose.yml new file mode 100644 index 0000000..c117858 --- /dev/null +++ b/.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 + diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ef5db3b --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +ollama_models +chroma_db_data \ No newline at end of file diff --git a/README.md b/README.md index 2250085..3862ce2 100644 --- a/README.md +++ b/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?" +``` \ No newline at end of file diff --git a/docker-compose-nvidia.yml b/docker-compose.nvidia.yml similarity index 100% rename from docker-compose-nvidia.yml rename to docker-compose.nvidia.yml diff --git a/docker-compose.yml b/docker-compose.yml index a588738..acc0c35 100644 --- a/docker-compose.yml +++ b/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 @@ -30,4 +30,7 @@ services: volumes: - ./ollama_models:/root/.ollama environment: - - OLLAMA_HOST=0.0.0.0 \ No newline at end of file + - OLLAMA_HOST=0.0.0.0 + + python-loader: + image: mcr.microsoft.com/devcontainers/base:jammy \ No newline at end of file diff --git a/flake.lock b/flake.lock deleted file mode 100644 index f4fb252..0000000 --- a/flake.lock +++ /dev/null @@ -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 -} diff --git a/flake.nix b/flake.nix deleted file mode 100644 index 304cf1b..0000000 --- a/flake.nix +++ /dev/null @@ -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 - ]; - }; - }; -} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..53cd720 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +chromadb +requests \ No newline at end of file diff --git a/script.py b/script.py new file mode 100644 index 0000000..fec5631 --- /dev/null +++ b/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) \ No newline at end of file