Skip to content

Commit 6718ed2

Browse files
authored
Feature image search (#5)
* Update our RAG to use CLIP embeddings to support the new multi-modl data. * Working reverse image search using our new multimodal DB, with good results for the current prompt * 1. Remove streaming endpoint 2. Actually pass images in to our vision model 3. Update code to use the latest 0.5.0 release of the vector db with separate image and text embeddings
1 parent e6a63fa commit 6718ed2

File tree

17 files changed

+324
-326
lines changed

17 files changed

+324
-326
lines changed

.env

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
DEFAULT_MODEL="qwen3:0.6b"
2+
DEFAULT_IMAGE_MODEL="qwen2.5vl:3b"
3+
24
OLLAMA_URL="http://ollama:11434"
35
DB_PATH="/data/enmemoryalpha_db"
46
COLLECTION_NAME="memoryalpha"

.github/workflows/ci-build.yml

Lines changed: 0 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -52,22 +52,6 @@ jobs:
5252
exit 1
5353
fi
5454
55-
- name: Test streaming endpoint
56-
run: |
57-
# Test the streaming endpoint
58-
response=$(timeout 30 curl -s -N -H "Accept: text/event-stream" \
59-
"http://localhost:8000/memoryalpha/rag/stream?question=What%20is%20a%20Transporter?&thinkingmode=DISABLED&max_tokens=50&top_k=3" \
60-
| head -10)
61-
62-
# Check if streaming response contains data events
63-
if echo "$response" | grep -q "data:"; then
64-
echo "✅ Streaming endpoint test passed"
65-
else
66-
echo "❌ Streaming endpoint test failed - no streaming data found"
67-
echo "Response: $response"
68-
exit 1
69-
fi
70-
7155
- name: Generate OpenAPI spec
7256
run: |
7357
# Download OpenAPI spec

Dockerfile

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,16 @@ FROM python:3.12-slim-bullseye AS devcontainer
33
ARG DEBIAN_FRONTEND=noninteractive
44

55
RUN apt-get update &&\
6-
apt-get -y install curl wget git
6+
apt-get -y install jq curl wget git
77

88
COPY ./requirements.txt /tmp/pip-tmp/
99
RUN pip install --no-cache-dir -r /tmp/pip-tmp/requirements.txt \
1010
&& rm -rf /tmp/pip-tmp
1111

1212
WORKDIR /data
13-
RUN wget https://github.com/aniongithub/memoryalpha-vectordb/releases/latest/download/enmemoryalpha_db.tar.gz &&\
13+
14+
ARG MEMORYALPHA_DB_RELEASE=v0.5.0
15+
RUN wget https://github.com/aniongithub/memoryalpha-vectordb/releases/download/${MEMORYALPHA_DB_RELEASE}/enmemoryalpha_db.tar.gz &&\
1416
tar -xzf enmemoryalpha_db.tar.gz &&\
1517
rm enmemoryalpha_db.tar.gz &&\
1618
chmod -R 0777 /data

README.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ A REST API for Retrieval-Augmented Generation (RAG) over Star Trek's MemoryAlpha
66

77
## Overview
88

9-
This project provides a streaming REST API that enables natural language queries over the comprehensive Star Trek MemoryAlpha database. It uses the vectorized database from [memoryalpha-vectordb](https://github.com/aniongithub/memoryalpha-vectordb) and combines it with local LLMs via Ollama to provide accurate, context-aware responses about Star Trek lore.
9+
This project provides a REST API that enables natural language queries over the comprehensive Star Trek MemoryAlpha database. It uses the vectorized database from [memoryalpha-vectordb](https://github.com/aniongithub/memoryalpha-vectordb) and combines it with local LLMs via Ollama to provide accurate, context-aware responses about Star Trek lore.
1010

1111
The system implements:
1212

@@ -169,7 +169,7 @@ graph TD
169169

170170
### Components
171171

172-
- **FastAPI**: REST API framework with streaming support
172+
- **FastAPI**: REST API framework and OpenAPI spec generation
173173
- **ChromaDB**: Vector database for document storage and retrieval
174174
- **Ollama**: Local LLM inference server
175175
- **Cross-Encoder**: Document reranking for improved relevance

api/main.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,8 @@
22
from contextlib import asynccontextmanager
33
from fastapi import FastAPI
44
from .memoryalpha.health import router as health_router
5-
from .memoryalpha.stream import router as stream_router
65
from .memoryalpha.ask import router as ask_router
6+
from .memoryalpha.identify import router as identify_router
77

88
# Configure logging
99
logging.basicConfig(level=logging.INFO)
@@ -21,5 +21,5 @@ async def lifespan(app: FastAPI):
2121
app = FastAPI(lifespan=lifespan)
2222

2323
app.include_router(health_router)
24-
app.include_router(stream_router)
25-
app.include_router(ask_router)
24+
app.include_router(ask_router)
25+
app.include_router(identify_router)

api/memoryalpha/identify.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
from fastapi import APIRouter, File, UploadFile, Query
2+
from fastapi.responses import JSONResponse
3+
import tempfile
4+
import os
5+
from .rag import MemoryAlphaRAG
6+
7+
router = APIRouter()
8+
9+
# Singleton or global instance for demo; in production, manage lifecycle properly
10+
rag_instance = MemoryAlphaRAG()
11+
12+
@router.post("/memoryalpha/rag/identify", summary="Multimodal Image Search")
13+
def identify_endpoint(
14+
file: UploadFile = File(...),
15+
top_k: int = Query(5, description="Number of results to return")
16+
):
17+
"""
18+
Accepts an image file upload, performs multimodal image search, and returns results.
19+
"""
20+
try:
21+
# Save uploaded file to a temp location
22+
with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[-1]) as tmp:
23+
tmp.write(file.file.read())
24+
image_path = tmp.name
25+
# Perform image search
26+
results = rag_instance.search_image(image_path, top_k=top_k)
27+
# Clean up temp file
28+
os.remove(image_path)
29+
return JSONResponse(content=results)
30+
except Exception as e:
31+
return JSONResponse(status_code=500, content={"error": str(e)})

0 commit comments

Comments
 (0)