-
Notifications
You must be signed in to change notification settings - Fork 259
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 all commits
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) | ||
|
Comment on lines
+117
to
+118
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This is really unfortunate if have to use this way to workaround the issue. What does it mean by "when FastAPI types are unavailable"? Missed a package? If add the package to dependency, will it pass the check? |
||
|
|
||
|
|
||
| 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]] | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. With this change, seems |
||
|
|
||
|
|
||
| @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)