7
7
from fastapi import FastAPI , Query
8
8
from fastapi .middleware .cors import CORSMiddleware
9
9
from fastapi .staticfiles import StaticFiles
10
+ from qdrant_client import QdrantClient
10
11
11
12
OLLAMA_URL = "http://localhost:11434/api/generate"
12
13
OLLAMA_MODEL = "gemma3"
14
+ QDRANT_GRPC_URL = os .getenv ("QDRANT_GRPC_URL" , "http://localhost:6334/" )
13
15
14
16
# 1. Extract caption from image using Ollama vision model
15
17
@cocoindex .op .function (cache = True , behavior_version = 1 )
@@ -42,7 +44,12 @@ def get_image_caption(img_bytes: bytes) -> str:
42
44
43
45
44
46
# 2. Embed the caption string
45
- def caption_to_embedding (caption : cocoindex .DataSlice ) -> cocoindex .DataSlice :
47
+ @cocoindex .transform_flow ()
48
+ def caption_to_embedding (caption : cocoindex .DataSlice [str ]) -> cocoindex .DataSlice [list [float ]]:
49
+ """
50
+ Embed the caption using a CLIP model.
51
+ This is shared logic between indexing and querying.
52
+ """
46
53
return caption .transform (
47
54
cocoindex .functions .SentenceTransformerEmbed (
48
55
model = "clip-ViT-L-14" ,
@@ -70,7 +77,7 @@ def image_object_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope:
70
77
"img_embeddings" ,
71
78
cocoindex .storages .Qdrant (
72
79
collection_name = "image_search" ,
73
- grpc_url = os . getenv ( " QDRANT_GRPC_URL" , "http://localhost:6334/" ) ,
80
+ grpc_url = QDRANT_GRPC_URL ,
74
81
),
75
82
primary_key_fields = ["id" ],
76
83
setup_by_user = True ,
@@ -93,26 +100,31 @@ def image_object_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope:
93
100
def startup_event ():
94
101
load_dotenv ()
95
102
cocoindex .init ()
96
- app .state .query_handler = cocoindex .query .SimpleSemanticsQueryHandler (
97
- name = "ImageObjectSearch" ,
98
- flow = image_object_embedding_flow ,
99
- target_name = "img_embeddings" ,
100
- query_transform_flow = caption_to_embedding ,
101
- default_similarity_metric = cocoindex .VectorSimilarityMetric .COSINE_SIMILARITY ,
103
+ # Initialize Qdrant client
104
+ app .state .qdrant_client = QdrantClient (
105
+ url = QDRANT_GRPC_URL ,
106
+ prefer_grpc = True
102
107
)
103
108
app .state .live_updater = cocoindex .FlowLiveUpdater (image_object_embedding_flow )
104
109
app .state .live_updater .start ()
105
110
106
111
@app .get ("/search" )
107
112
def search (q : str = Query (..., description = "Search query" ), limit : int = Query (5 , description = "Number of results" )):
108
- query_handler = app .state .query_handler
109
- results , _ = query_handler .search (q , limit , "embedding" )
113
+ # Get the embedding for the query
114
+ query_embedding = caption_to_embedding .eval (q )
115
+
116
+ # Search in Qdrant
117
+ search_results = app .state .qdrant_client .search (
118
+ collection_name = "image_search" ,
119
+ query_vector = ("embedding" , query_embedding ),
120
+ limit = limit
121
+ )
122
+
123
+ # Format results
110
124
out = []
111
- for result in results :
112
- row = dict (result .data )
113
- # Only include filename and score
125
+ for result in search_results :
114
126
out .append ({
115
- "filename" : row ["filename" ],
127
+ "filename" : result . payload ["filename" ],
116
128
"score" : result .score
117
129
})
118
130
return {"results" : out }
0 commit comments