-
Notifications
You must be signed in to change notification settings - Fork 267
examples: fix mypy errors in example entry points; add local type-check script (#1091) #1248
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 1 commit
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,34 @@ | ||
| $ErrorActionPreference = 'Stop' | ||
|
|
||
| # Resolve python in local venv | ||
| $repoRoot = Split-Path -Parent $PSScriptRoot | ||
| $python = Join-Path $repoRoot '.venv\Scripts\python.exe' | ||
| if (-not (Test-Path $python)) { | ||
| $python = 'python' | ||
| } | ||
|
|
||
| # Ensure mypy can resolve local cocoindex package sources | ||
| $env:MYPYPATH = Join-Path $repoRoot 'python' | ||
|
|
||
| # Collect example entry files | ||
| $examples = Join-Path $repoRoot 'examples' | ||
| $files = Get-ChildItem -Path $examples -Recurse -File | | ||
| Where-Object { $_.Name -in @('main.py','colpali_main.py') } | | ||
| Sort-Object FullName | ||
|
|
||
| $failed = @() | ||
| foreach ($f in $files) { | ||
| Write-Host (">>> Checking " + $f.FullName) | ||
| & $python -m mypy --ignore-missing-imports --follow-imports=silent $f.FullName | ||
| if ($LASTEXITCODE -ne 0) { | ||
| $failed += $f.FullName | ||
| } | ||
| } | ||
|
|
||
| if ($failed.Count -gt 0) { | ||
| Write-Host "\nFailures:" | ||
| $failed | ForEach-Object { Write-Host $_ } | ||
| exit 1 | ||
| } else { | ||
| Write-Host "\nAll example entry files passed mypy." | ||
| } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -6,6 +6,7 @@ | |
| from psycopg_pool import ConnectionPool | ||
| from contextlib import asynccontextmanager | ||
| import os | ||
| from typing import Any, AsyncIterator | ||
|
|
||
|
|
||
| @cocoindex.transform_flow() | ||
|
|
@@ -26,7 +27,7 @@ def text_to_embedding( | |
| @cocoindex.flow_def(name="MarkdownEmbeddingFastApiExample") | ||
| def markdown_embedding_flow( | ||
| flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope | ||
| ): | ||
| ) -> None: | ||
| """ | ||
| Define an example flow that embeds markdown files into a vector database. | ||
| """ | ||
|
|
@@ -65,7 +66,7 @@ def markdown_embedding_flow( | |
| ) | ||
|
|
||
|
|
||
| def search(pool: ConnectionPool, query: str, top_k: int = 5): | ||
| def search(pool: ConnectionPool, query: str, top_k: int = 5) -> list[dict[str, Any]]: | ||
| # Get the table name, for the export target in the text_embedding_flow above. | ||
| table_name = cocoindex.utils.get_target_default_name( | ||
| markdown_embedding_flow, "doc_embeddings" | ||
|
|
@@ -89,7 +90,7 @@ def search(pool: ConnectionPool, query: str, top_k: int = 5): | |
|
|
||
|
|
||
| @asynccontextmanager | ||
| def lifespan(app: FastAPI): | ||
| async def lifespan(app: FastAPI) -> AsyncIterator[None]: | ||
| load_dotenv() | ||
| cocoindex.init() | ||
| pool = ConnectionPool(os.getenv("COCOINDEX_DATABASE_URL")) | ||
|
|
@@ -103,16 +104,19 @@ def lifespan(app: FastAPI): | |
| fastapi_app = FastAPI(lifespan=lifespan) | ||
|
|
||
|
|
||
| @fastapi_app.get("/search") | ||
| def search_endpoint( | ||
| request: Request, | ||
| q: str = Query(..., description="Search query"), | ||
| limit: int = Query(5, description="Number of results"), | ||
| ): | ||
| ) -> dict[str, Any]: | ||
| pool = request.app.state.pool | ||
| results = search(pool, q, limit) | ||
| return {"results": results} | ||
|
|
||
|
|
||
| # Attach route without using decorator to avoid untyped-decorator when FastAPI types are unavailable | ||
| fastapi_app.get("/search")(search_endpoint) | ||
|
||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| uvicorn.run(fastapi_app, host="0.0.0.0", port=8080) | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -3,7 +3,7 @@ | |
| import io | ||
| import os | ||
| from contextlib import asynccontextmanager | ||
| from typing import Any, Literal | ||
| from typing import Any, Literal, Final, TypeAlias, cast, AsyncIterator | ||
|
|
||
| import cocoindex | ||
| import torch | ||
|
|
@@ -19,7 +19,8 @@ | |
| QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6334/") | ||
| QDRANT_COLLECTION = "ImageSearch" | ||
| CLIP_MODEL_NAME = "openai/clip-vit-large-patch14" | ||
| CLIP_MODEL_DIMENSION = 768 | ||
| CLIP_MODEL_DIMENSION: Final[int] = 768 | ||
| CLIPVector: TypeAlias = cocoindex.Vector[cocoindex.Float32, Literal[768]] | ||
|
||
|
|
||
|
|
||
| @functools.cache | ||
|
|
@@ -37,13 +38,13 @@ def embed_query(text: str) -> list[float]: | |
| inputs = processor(text=[text], return_tensors="pt", padding=True) | ||
| with torch.no_grad(): | ||
| features = model.get_text_features(**inputs) | ||
| return features[0].tolist() | ||
| return cast(list[float], features[0].tolist()) | ||
|
|
||
|
|
||
| @cocoindex.op.function(cache=True, behavior_version=1, gpu=True) | ||
| def embed_image( | ||
| img_bytes: bytes, | ||
| ) -> cocoindex.Vector[cocoindex.Float32, Literal[CLIP_MODEL_DIMENSION]]: | ||
| ) -> CLIPVector: | ||
| """ | ||
| Convert image to embedding using CLIP model. | ||
| """ | ||
|
|
@@ -52,7 +53,7 @@ def embed_image( | |
| inputs = processor(images=image, return_tensors="pt") | ||
| with torch.no_grad(): | ||
| features = model.get_image_features(**inputs) | ||
| return features[0].tolist() | ||
| return cast(CLIPVector, features[0].tolist()) | ||
|
|
||
|
|
||
| # CocoIndex flow: Ingest images, extract captions, embed, export to Qdrant | ||
|
|
@@ -112,7 +113,7 @@ def image_object_embedding_flow( | |
|
|
||
|
|
||
| @asynccontextmanager | ||
| async def lifespan(app: FastAPI) -> None: | ||
| async def lifespan(app: FastAPI) -> AsyncIterator[None]: | ||
| load_dotenv() | ||
| cocoindex.init() | ||
| image_object_embedding_flow.setup(report_to_stdout=True) | ||
|
|
@@ -141,11 +142,10 @@ async def lifespan(app: FastAPI) -> None: | |
|
|
||
|
|
||
| # --- Search API --- | ||
| @app.get("/search") | ||
| def search( | ||
| q: str = Query(..., description="Search query"), | ||
| limit: int = Query(5, description="Number of results"), | ||
| ) -> Any: | ||
| ) -> dict[str, Any]: | ||
| # Get the embedding for the query | ||
| query_embedding = embed_query(q) | ||
|
|
||
|
|
@@ -169,3 +169,7 @@ def search( | |
| for result in search_results | ||
| ] | ||
| } | ||
|
|
||
|
|
||
| # Attach route without using decorator to avoid untyped-decorator when FastAPI types are unavailable | ||
| app.get("/search")(search) | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm thinking about also considering code simplicity/readability. Given these examples are for users to understand how to use cocoindex, this matters.
For case like this, it's essentially the underlying library doesn't have a specific type, and the return type of the current function is clear.
I think we can just ignore it by a comment like
# type: ignore(IMO
cast(...)is more useful for values within a function - from the point on, the type is clear)