33import io
44import os
55from contextlib import asynccontextmanager
6- from typing import Any , Literal
6+ from typing import Any , Literal , Final , TypeAlias , cast , AsyncIterator
77
88import cocoindex
99import torch
1919QDRANT_URL = os .getenv ("QDRANT_URL" , "http://localhost:6334/" )
2020QDRANT_COLLECTION = "ImageSearch"
2121CLIP_MODEL_NAME = "openai/clip-vit-large-patch14"
22- CLIP_MODEL_DIMENSION = 768
22+ CLIP_MODEL_DIMENSION : Final [int ] = 768
23+ CLIPVector : TypeAlias = cocoindex .Vector [cocoindex .Float32 , Literal [768 ]]
2324
2425
2526@functools .cache
@@ -37,13 +38,13 @@ def embed_query(text: str) -> list[float]:
3738 inputs = processor (text = [text ], return_tensors = "pt" , padding = True )
3839 with torch .no_grad ():
3940 features = model .get_text_features (** inputs )
40- return features [0 ].tolist ()
41+ return cast ( list [ float ], features [0 ].tolist () )
4142
4243
4344@cocoindex .op .function (cache = True , behavior_version = 1 , gpu = True )
4445def embed_image (
4546 img_bytes : bytes ,
46- ) -> cocoindex . Vector [ cocoindex . Float32 , Literal [ CLIP_MODEL_DIMENSION ]] :
47+ ) -> CLIPVector :
4748 """
4849 Convert image to embedding using CLIP model.
4950 """
@@ -52,7 +53,7 @@ def embed_image(
5253 inputs = processor (images = image , return_tensors = "pt" )
5354 with torch .no_grad ():
5455 features = model .get_image_features (** inputs )
55- return features [0 ].tolist ()
56+ return cast ( CLIPVector , features [0 ].tolist () )
5657
5758
5859# CocoIndex flow: Ingest images, extract captions, embed, export to Qdrant
@@ -112,7 +113,7 @@ def image_object_embedding_flow(
112113
113114
114115@asynccontextmanager
115- async def lifespan (app : FastAPI ) -> None :
116+ async def lifespan (app : FastAPI ) -> AsyncIterator [ None ] :
116117 load_dotenv ()
117118 cocoindex .init ()
118119 image_object_embedding_flow .setup (report_to_stdout = True )
@@ -141,11 +142,10 @@ async def lifespan(app: FastAPI) -> None:
141142
142143
143144# --- Search API ---
144- @app .get ("/search" )
145145def search (
146146 q : str = Query (..., description = "Search query" ),
147147 limit : int = Query (5 , description = "Number of results" ),
148- ) -> Any :
148+ ) -> dict [ str , Any ] :
149149 # Get the embedding for the query
150150 query_embedding = embed_query (q )
151151
@@ -169,3 +169,7 @@ def search(
169169 for result in search_results
170170 ]
171171 }
172+
173+
174+ # Attach route without using decorator to avoid untyped-decorator when FastAPI types are unavailable
175+ app .get ("/search" )(search )
0 commit comments