Make potentially blocking calls async

Author: sunu

pyproject.toml

@@ -18,6 +18,7 @@ dependencies = [
     "openai",
     "pydantic-ai",
     "shapely",
+    "aiohttp",
 ]
 
 [tool.setuptools]

stac_search/agents/collections_search.py

@@ -2,6 +2,7 @@
 Search module for STAC Natural Query
 """
 
+import asyncio
 import logging
 import os
 import time
@@ -97,10 +98,11 @@ async def collection_search(
     logger.info(f"Model loading time: {load_model_time - start_time:.4f} seconds")
 
     # Generate query embedding
-    query_embedding = catalog_manager.model.encode([query])
+    query_embedding = await asyncio.to_thread(catalog_manager.model.encode, [query])
 
     # Search vector database
-    results = collection.query(
+    results = await asyncio.to_thread(
+        collection.query,
         query_embeddings=query_embedding.tolist(),
         n_results=top_k * 2,  # Get more results initially for better reranking
     )
@@ -131,6 +133,4 @@ async def main():
 
 
 if __name__ == "__main__":
-    import asyncio
-
     asyncio.run(main())

stac_search/agents/items_search.py

@@ -5,9 +5,10 @@
 from dataclasses import dataclass
 from pprint import pformat
 import time
+import asyncio
 from typing import List, Dict, Any, Union
 
-import requests
+import aiohttp
 from pydantic_ai import Agent, RunContext
 from pystac_client import Client
 from pydantic import BaseModel, ConfigDict
@@ -259,15 +260,13 @@ async def construct_cql2_filter(ctx: RunContext[Context]) -> FilterExpr | None:
     return await cql2_filter_agent.run(ctx.deps.query)
 
 
-def get_polygon_from_geodini(location: str):
+async def get_polygon_from_geodini(location: str):
     geodini_api = f"{GEODINI_API}/search_complex"
-    response = requests.get(
-        geodini_api,
-        params={"query": location},
-    )
-    result = response.json().get("result", None)
-    if result:
-        return result.get("geometry", None)
+    async with aiohttp.ClientSession() as session:
+        async with session.get(geodini_api, params={"query": location}) as response:
+            result = (await response.json()).get("result", None)
+            if result:
+                return result.get("geometry", None)
     return None
 
 
@@ -300,10 +299,12 @@ async def item_search(ctx: Context) -> ItemSearchResult:
         # If no specific collections were found, use the default target collections
         default_target_collections = DEFAULT_TARGET_COLLECTIONS
         # check that default_target_collections exist in the catalog
-        all_collection_ids = [
-            collection.id
-            for collection in Client.open(catalog_url_to_use).get_collections()
-        ]
+        all_collection_ids = await asyncio.to_thread(
+            lambda: [
+                collection.id
+                for collection in Client.open(catalog_url_to_use).get_collections()
+            ]
+        )
         default_target_collections = [
             collection_id
             for collection_id in default_target_collections
@@ -339,7 +340,7 @@ async def item_search(ctx: Context) -> ItemSearchResult:
         f"Params formulation time: {params_formulation_time - query_formulation_time} seconds"
     )
 
-    polygon = get_polygon_from_geodini(results.data.location)
+    polygon = await get_polygon_from_geodini(results.data.location)
     if polygon:
         logger.info(f"Found polygon for {results.data.location}")
         params["intersects"] = polygon
@@ -349,25 +350,21 @@ async def item_search(ctx: Context) -> ItemSearchResult:
             items=None, search_params=params, aoi=None, explanation=explanation
         )
     geocoding_time = time.time()
-    logger.info(
-        f"Geocoding time: {geocoding_time - params_formulation_time} seconds"
-    )
+    logger.info(f"Geocoding time: {geocoding_time - params_formulation_time} seconds")
 
     if ctx.return_search_params_only:
         logger.info("Returning STAC query parameters only")
         total_time = time.time() - start_time
-        logger.info(
-            f"Total time: {total_time} seconds"
-        )
+        logger.info(f"Total time: {total_time} seconds")
         return ItemSearchResult(
             search_params=params, aoi=polygon, explanation=explanation
         )
 
-    items = list(client.search(**params).items_as_dicts())
-    search_time = time.time()
-    logger.info(
-        f"Search time: {search_time - geocoding_time} seconds"
+    items = await asyncio.to_thread(
+        lambda: list(client.search(**params).items_as_dicts())
     )
+    search_time = time.time()
+    logger.info(f"Search time: {search_time - geocoding_time} seconds")
     total_time = time.time() - start_time
     logger.info(f"Total time: {total_time} seconds")
     return ItemSearchResult(
@@ -382,6 +379,4 @@ async def main():
 
 
 if __name__ == "__main__":
-    import asyncio
-
     asyncio.run(main())

stac_search/catalog_manager.py

@@ -2,6 +2,7 @@
 Catalog Manager for STAC Natural Query - handles dynamic catalog loading and management
 """
 
+import asyncio
 import hashlib
 import logging
 import os
@@ -52,38 +53,46 @@ def catalog_exists(self, catalog_url: str) -> bool:
             logger.error(f"Error checking catalog existence: {e}")
             return False
 
-    def validate_catalog_url(self, catalog_url: str) -> bool:
+    async def validate_catalog_url(self, catalog_url: str) -> bool:
         """Validate that the catalog URL is accessible and is a valid STAC catalog"""
         try:
-            stac_client = Client.open(catalog_url)
-            # Try to get at least one collection to verify it's a valid catalog
-            collections = list(stac_client.collection_search().collections())
-            return len(collections) > 0
+
+            def _validate():
+                stac_client = Client.open(catalog_url)
+                # Try to get at least one collection to verify it's a valid catalog
+                collections = list(stac_client.collection_search().collections())
+                return len(collections) > 0
+
+            return await asyncio.to_thread(_validate)
         except Exception as e:
             logger.error(f"Invalid catalog URL {catalog_url}: {e}")
             return False
 
-    def fetch_collections(self, stac_client: Client) -> list:
+    async def fetch_collections(self, stac_client: Client) -> list:
         """Fetch STAC collections using pystac-client"""
         try:
-            collections = stac_client.collection_search().collections()
-            return list(collections)
+
+            def _fetch():
+                collections = stac_client.collection_search().collections()
+                return list(collections)
+
+            return await asyncio.to_thread(_fetch)
         except Exception as e:
             logger.error(f"Error fetching collections: {e}")
             return []
 
-    def generate_embeddings(self, collections: list) -> list:
+    async def generate_embeddings(self, collections: list) -> list:
         """Generate embeddings for each collection (title + description)"""
         texts = []
         for collection in collections:
             title = getattr(collection, "title", "") or ""
             description = getattr(collection, "description", "") or ""
             texts.append(f"{title} {description}")
 
-        embeddings = self.model.encode(texts)
+        embeddings = await asyncio.to_thread(self.model.encode, texts)
         return embeddings
 
-    def store_in_vector_db(self, collections: list, chroma_collection) -> None:
+    async def store_in_vector_db(self, collections: list, chroma_collection) -> None:
         """Store embeddings in ChromaDB"""
         if not collections:
             logger.warning("No collections to store")
@@ -98,9 +107,10 @@ def store_in_vector_db(self, collections: list, chroma_collection) -> None:
             }
             metadatas.append(metadata)
 
-        embeddings = self.generate_embeddings(collections)
+        embeddings = await self.generate_embeddings(collections)
 
-        chroma_collection.add(
+        await asyncio.to_thread(
+            chroma_collection.add,
             ids=[str(i) for i in range(len(collections))],
             embeddings=embeddings,
             metadatas=metadatas,
@@ -110,7 +120,7 @@ async def load_catalog(self, catalog_url: str) -> Dict[str, Any]:
         """Load and index a catalog if it doesn't exist"""
         try:
             # Validate catalog URL first
-            if not self.validate_catalog_url(catalog_url):
+            if not await self.validate_catalog_url(catalog_url):
                 return {
                     "success": False,
                     "error": f"Invalid or inaccessible catalog URL: {catalog_url}",
@@ -127,8 +137,8 @@ async def load_catalog(self, catalog_url: str) -> Dict[str, Any]:
 
             # Load the catalog
             logger.info(f"Loading catalog from {catalog_url}")
-            stac_client = Client.open(catalog_url)
-            collections = self.fetch_collections(stac_client)
+            stac_client = await asyncio.to_thread(Client.open, catalog_url)
+            collections = await self.fetch_collections(stac_client)
 
             if not collections:
                 return {
@@ -143,7 +153,7 @@ async def load_catalog(self, catalog_url: str) -> Dict[str, Any]:
             )
 
             # Store in vector database
-            self.store_in_vector_db(collections, chroma_collection)
+            await self.store_in_vector_db(collections, chroma_collection)
 
             logger.info(
                 f"Successfully indexed {len(collections)} collections from {catalog_url}"