weave-python 0.28.2__py3-none-any.whl → 0.30.3__py3-none-any.whl

weave/weavesql/llmxdb/models.py

@@ -0,0 +1,594 @@
+# Code generated by sqlc. DO NOT EDIT.
+# versions:
+#   sqlc v1.28.0
+# source: models.sql
+import dataclasses
+import datetime
+from typing import Any, AsyncIterator, Iterator, List, Optional
+import uuid
+
+import sqlalchemy
+import sqlalchemy.ext.asyncio
+
+from weave.weavesql.llmxdb import models
+
+
+BATCH_UPSERT_MODELS = """-- name: batch_upsert_models \\:copyfrom
+INSERT INTO llm_models (provider_id, model_id, slug, name, display_name, description, version,
+                        classification, architecture, capabilities, performance, token_info,
+                        pricing, api_details, training, licensing, safety, availability,
+                        technical_specs, metadata)
+VALUES (:p1, :p2, :p3, :p4, :p5, :p6, :p7,
+        :p8, :p9, :p10, :p11, :p12,
+        :p13, :p14, :p15, :p16, :p17, :p18,
+        :p19, :p20)
+"""
+
+
+@dataclasses.dataclass()
+class BatchUpsertModelsParams:
+    provider_id: uuid.UUID
+    model_id: str
+    slug: str
+    name: str
+    display_name: Optional[str]
+    description: Optional[str]
+    version: Optional[str]
+    classification: Any
+    architecture: Any
+    capabilities: Any
+    performance: Any
+    token_info: Any
+    pricing: Any
+    api_details: Any
+    training: Any
+    licensing: Any
+    safety: Any
+    availability: Any
+    technical_specs: Any
+    metadata: Any
+
+
+GET_MODEL = """-- name: get_model \\:one
+SELECT id, provider_id, model_id, slug, name, display_name, description, version, capabilities, classification, architecture, performance, token_info, pricing, api_details, training, licensing, safety, availability, technical_specs, release_date, training_data_cutoff, deprecation_date, metadata, last_scraped_at, is_active, is_deprecated, created_at, updated_at
+FROM llm_models
+WHERE slug = :p1
+  AND is_active = true
+"""
+
+
+GET_MODELS = """-- name: get_models \\:many
+SELECT id, provider_id, model_id, slug, name, display_name, description, version, capabilities, classification, architecture, performance, token_info, pricing, api_details, training, licensing, safety, availability, technical_specs, release_date, training_data_cutoff, deprecation_date, metadata, last_scraped_at, is_active, is_deprecated, created_at, updated_at
+FROM llm_models
+WHERE is_active = true
+ORDER BY provider_id, name
+"""
+
+
+GET_MODELS_BY_PROVIDER = """-- name: get_models_by_provider \\:many
+SELECT id, provider_id, model_id, slug, name, display_name, description, version, capabilities, classification, architecture, performance, token_info, pricing, api_details, training, licensing, safety, availability, technical_specs, release_date, training_data_cutoff, deprecation_date, metadata, last_scraped_at, is_active, is_deprecated, created_at, updated_at
+FROM llm_models
+WHERE provider_id = :p1
+  AND is_active = true
+ORDER BY name
+"""
+
+
+UPDATE_MODEL_DEPRECATION = """-- name: update_model_deprecation \\:exec
+UPDATE llm_models
+SET is_deprecated    = true,
+    deprecation_date = :p1,
+    metadata         = jsonb_set(metadata, '{replacement_model_id}', to_jsonb(:p2\\:\\:text))
+WHERE slug = :p3
+"""
+
+
+@dataclasses.dataclass()
+class UpdateModelDeprecationParams:
+    deprecation_date: Optional[datetime.datetime]
+    replacement_model_id: str
+    slug: str
+
+
+UPSERT_MODEL = """-- name: upsert_model \\:one
+INSERT INTO llm_models (provider_id, model_id, slug, name, display_name, description, version,
+                        classification, architecture, capabilities, performance, token_info,
+                        pricing, api_details, training, licensing, safety, availability,
+                        technical_specs, release_date, training_data_cutoff, deprecation_date,
+                        metadata, last_scraped_at)
+VALUES (:p1, :p2, :p3, :p4, :p5, :p6, :p7,
+        :p8, :p9, :p10, :p11, :p12,
+        :p13, :p14, :p15, :p16, :p17, :p18,
+        :p19, :p20, :p21, :p22,
+        :p23, NOW())
+ON CONFLICT (slug) DO UPDATE SET name                 = EXCLUDED.name,
+                                 display_name         = EXCLUDED.display_name,
+                                 description          = EXCLUDED.description,
+                                 version              = EXCLUDED.version,
+                                 classification       = EXCLUDED.classification,
+                                 architecture         = EXCLUDED.architecture,
+                                 capabilities         = EXCLUDED.capabilities,
+                                 performance          = EXCLUDED.performance,
+                                 token_info           = EXCLUDED.token_info,
+                                 pricing              = EXCLUDED.pricing,
+                                 api_details          = EXCLUDED.api_details,
+                                 training             = EXCLUDED.training,
+                                 licensing            = EXCLUDED.licensing,
+                                 safety               = EXCLUDED.safety,
+                                 availability         = EXCLUDED.availability,
+                                 technical_specs      = EXCLUDED.technical_specs,
+                                 release_date         = EXCLUDED.release_date,
+                                 training_data_cutoff = EXCLUDED.training_data_cutoff,
+                                 deprecation_date     = EXCLUDED.deprecation_date,
+                                 metadata             = EXCLUDED.metadata,
+                                 last_scraped_at      = NOW(),
+                                 updated_at           = NOW()
+RETURNING id, provider_id, model_id, slug, name, display_name, description, version, capabilities, classification, architecture, performance, token_info, pricing, api_details, training, licensing, safety, availability, technical_specs, release_date, training_data_cutoff, deprecation_date, metadata, last_scraped_at, is_active, is_deprecated, created_at, updated_at
+"""
+
+
+@dataclasses.dataclass()
+class UpsertModelParams:
+    provider_id: uuid.UUID
+    model_id: str
+    slug: str
+    name: str
+    display_name: Optional[str]
+    description: Optional[str]
+    version: Optional[str]
+    classification: Any
+    architecture: Any
+    capabilities: Any
+    performance: Any
+    token_info: Any
+    pricing: Any
+    api_details: Any
+    training: Any
+    licensing: Any
+    safety: Any
+    availability: Any
+    technical_specs: Any
+    release_date: Optional[datetime.datetime]
+    training_data_cutoff: Optional[datetime.datetime]
+    deprecation_date: Optional[datetime.datetime]
+    metadata: Any
+
+
+class Querier:
+    def __init__(self, conn: sqlalchemy.engine.Connection):
+        self._conn = conn
+
+    def batch_upsert_models(self, arg_list: List[BatchUpsertModelsParams]) -> int:
+        data = list()
+        for item in arg_list:
+            data.append(
+                {
+                    "p1": item.provider_id,
+                    "p2": item.model_id,
+                    "p3": item.slug,
+                    "p4": item.name,
+                    "p5": item.display_name,
+                    "p6": item.description,
+                    "p7": item.version,
+                    "p8": item.classification,
+                    "p9": item.architecture,
+                    "p10": item.capabilities,
+                    "p11": item.performance,
+                    "p12": item.token_info,
+                    "p13": item.pricing,
+                    "p14": item.api_details,
+                    "p15": item.training,
+                    "p16": item.licensing,
+                    "p17": item.safety,
+                    "p18": item.availability,
+                    "p19": item.technical_specs,
+                    "p20": item.metadata,
+                }
+            )
+        result = self._conn.execute(sqlalchemy.text(BATCH_UPSERT_MODELS), data)
+        return result.rowcount
+
+    def get_model(self, *, slug: str) -> Optional[models.LlmModel]:
+        row = self._conn.execute(sqlalchemy.text(GET_MODEL), {"p1": slug}).first()
+        if row is None:
+            return None
+        return models.LlmModel(
+            id=row[0],
+            provider_id=row[1],
+            model_id=row[2],
+            slug=row[3],
+            name=row[4],
+            display_name=row[5],
+            description=row[6],
+            version=row[7],
+            capabilities=row[8],
+            classification=row[9],
+            architecture=row[10],
+            performance=row[11],
+            token_info=row[12],
+            pricing=row[13],
+            api_details=row[14],
+            training=row[15],
+            licensing=row[16],
+            safety=row[17],
+            availability=row[18],
+            technical_specs=row[19],
+            release_date=row[20],
+            training_data_cutoff=row[21],
+            deprecation_date=row[22],
+            metadata=row[23],
+            last_scraped_at=row[24],
+            is_active=row[25],
+            is_deprecated=row[26],
+            created_at=row[27],
+            updated_at=row[28],
+        )
+
+    def get_models(self) -> Iterator[models.LlmModel]:
+        result = self._conn.execute(sqlalchemy.text(GET_MODELS))
+        for row in result:
+            yield models.LlmModel(
+                id=row[0],
+                provider_id=row[1],
+                model_id=row[2],
+                slug=row[3],
+                name=row[4],
+                display_name=row[5],
+                description=row[6],
+                version=row[7],
+                capabilities=row[8],
+                classification=row[9],
+                architecture=row[10],
+                performance=row[11],
+                token_info=row[12],
+                pricing=row[13],
+                api_details=row[14],
+                training=row[15],
+                licensing=row[16],
+                safety=row[17],
+                availability=row[18],
+                technical_specs=row[19],
+                release_date=row[20],
+                training_data_cutoff=row[21],
+                deprecation_date=row[22],
+                metadata=row[23],
+                last_scraped_at=row[24],
+                is_active=row[25],
+                is_deprecated=row[26],
+                created_at=row[27],
+                updated_at=row[28],
+            )
+
+    def get_models_by_provider(
+        self, *, provider_id: uuid.UUID
+    ) -> Iterator[models.LlmModel]:
+        result = self._conn.execute(
+            sqlalchemy.text(GET_MODELS_BY_PROVIDER), {"p1": provider_id}
+        )
+        for row in result:
+            yield models.LlmModel(
+                id=row[0],
+                provider_id=row[1],
+                model_id=row[2],
+                slug=row[3],
+                name=row[4],
+                display_name=row[5],
+                description=row[6],
+                version=row[7],
+                capabilities=row[8],
+                classification=row[9],
+                architecture=row[10],
+                performance=row[11],
+                token_info=row[12],
+                pricing=row[13],
+                api_details=row[14],
+                training=row[15],
+                licensing=row[16],
+                safety=row[17],
+                availability=row[18],
+                technical_specs=row[19],
+                release_date=row[20],
+                training_data_cutoff=row[21],
+                deprecation_date=row[22],
+                metadata=row[23],
+                last_scraped_at=row[24],
+                is_active=row[25],
+                is_deprecated=row[26],
+                created_at=row[27],
+                updated_at=row[28],
+            )
+
+    def update_model_deprecation(self, arg: UpdateModelDeprecationParams) -> None:
+        self._conn.execute(
+            sqlalchemy.text(UPDATE_MODEL_DEPRECATION),
+            {
+                "p1": arg.deprecation_date,
+                "p2": arg.replacement_model_id,
+                "p3": arg.slug,
+            },
+        )
+
+    def upsert_model(self, arg: UpsertModelParams) -> Optional[models.LlmModel]:
+        row = self._conn.execute(
+            sqlalchemy.text(UPSERT_MODEL),
+            {
+                "p1": arg.provider_id,
+                "p2": arg.model_id,
+                "p3": arg.slug,
+                "p4": arg.name,
+                "p5": arg.display_name,
+                "p6": arg.description,
+                "p7": arg.version,
+                "p8": arg.classification,
+                "p9": arg.architecture,
+                "p10": arg.capabilities,
+                "p11": arg.performance,
+                "p12": arg.token_info,
+                "p13": arg.pricing,
+                "p14": arg.api_details,
+                "p15": arg.training,
+                "p16": arg.licensing,
+                "p17": arg.safety,
+                "p18": arg.availability,
+                "p19": arg.technical_specs,
+                "p20": arg.release_date,
+                "p21": arg.training_data_cutoff,
+                "p22": arg.deprecation_date,
+                "p23": arg.metadata,
+            },
+        ).first()
+        if row is None:
+            return None
+        return models.LlmModel(
+            id=row[0],
+            provider_id=row[1],
+            model_id=row[2],
+            slug=row[3],
+            name=row[4],
+            display_name=row[5],
+            description=row[6],
+            version=row[7],
+            capabilities=row[8],
+            classification=row[9],
+            architecture=row[10],
+            performance=row[11],
+            token_info=row[12],
+            pricing=row[13],
+            api_details=row[14],
+            training=row[15],
+            licensing=row[16],
+            safety=row[17],
+            availability=row[18],
+            technical_specs=row[19],
+            release_date=row[20],
+            training_data_cutoff=row[21],
+            deprecation_date=row[22],
+            metadata=row[23],
+            last_scraped_at=row[24],
+            is_active=row[25],
+            is_deprecated=row[26],
+            created_at=row[27],
+            updated_at=row[28],
+        )
+
+
+class AsyncQuerier:
+    def __init__(self, conn: sqlalchemy.ext.asyncio.AsyncConnection):
+        self._conn = conn
+
+    async def batch_upsert_models(self, arg_list: List[BatchUpsertModelsParams]) -> int:
+        data = list()
+        for item in arg_list:
+            data.append(
+                {
+                    "p1": item.provider_id,
+                    "p2": item.model_id,
+                    "p3": item.slug,
+                    "p4": item.name,
+                    "p5": item.display_name,
+                    "p6": item.description,
+                    "p7": item.version,
+                    "p8": item.classification,
+                    "p9": item.architecture,
+                    "p10": item.capabilities,
+                    "p11": item.performance,
+                    "p12": item.token_info,
+                    "p13": item.pricing,
+                    "p14": item.api_details,
+                    "p15": item.training,
+                    "p16": item.licensing,
+                    "p17": item.safety,
+                    "p18": item.availability,
+                    "p19": item.technical_specs,
+                    "p20": item.metadata,
+                }
+            )
+        result = await self._conn.execute(sqlalchemy.text(BATCH_UPSERT_MODELS), data)
+        return result.rowcount
+
+    async def get_model(self, *, slug: str) -> Optional[models.LlmModel]:
+        row = (
+            await self._conn.execute(sqlalchemy.text(GET_MODEL), {"p1": slug})
+        ).first()
+        if row is None:
+            return None
+        return models.LlmModel(
+            id=row[0],
+            provider_id=row[1],
+            model_id=row[2],
+            slug=row[3],
+            name=row[4],
+            display_name=row[5],
+            description=row[6],
+            version=row[7],
+            capabilities=row[8],
+            classification=row[9],
+            architecture=row[10],
+            performance=row[11],
+            token_info=row[12],
+            pricing=row[13],
+            api_details=row[14],
+            training=row[15],
+            licensing=row[16],
+            safety=row[17],
+            availability=row[18],
+            technical_specs=row[19],
+            release_date=row[20],
+            training_data_cutoff=row[21],
+            deprecation_date=row[22],
+            metadata=row[23],
+            last_scraped_at=row[24],
+            is_active=row[25],
+            is_deprecated=row[26],
+            created_at=row[27],
+            updated_at=row[28],
+        )
+
+    async def get_models(self) -> AsyncIterator[models.LlmModel]:
+        result = await self._conn.stream(sqlalchemy.text(GET_MODELS))
+        async for row in result:
+            yield models.LlmModel(
+                id=row[0],
+                provider_id=row[1],
+                model_id=row[2],
+                slug=row[3],
+                name=row[4],
+                display_name=row[5],
+                description=row[6],
+                version=row[7],
+                capabilities=row[8],
+                classification=row[9],
+                architecture=row[10],
+                performance=row[11],
+                token_info=row[12],
+                pricing=row[13],
+                api_details=row[14],
+                training=row[15],
+                licensing=row[16],
+                safety=row[17],
+                availability=row[18],
+                technical_specs=row[19],
+                release_date=row[20],
+                training_data_cutoff=row[21],
+                deprecation_date=row[22],
+                metadata=row[23],
+                last_scraped_at=row[24],
+                is_active=row[25],
+                is_deprecated=row[26],
+                created_at=row[27],
+                updated_at=row[28],
+            )
+
+    async def get_models_by_provider(
+        self, *, provider_id: uuid.UUID
+    ) -> AsyncIterator[models.LlmModel]:
+        result = await self._conn.stream(
+            sqlalchemy.text(GET_MODELS_BY_PROVIDER), {"p1": provider_id}
+        )
+        async for row in result:
+            yield models.LlmModel(
+                id=row[0],
+                provider_id=row[1],
+                model_id=row[2],
+                slug=row[3],
+                name=row[4],
+                display_name=row[5],
+                description=row[6],
+                version=row[7],
+                capabilities=row[8],
+                classification=row[9],
+                architecture=row[10],
+                performance=row[11],
+                token_info=row[12],
+                pricing=row[13],
+                api_details=row[14],
+                training=row[15],
+                licensing=row[16],
+                safety=row[17],
+                availability=row[18],
+                technical_specs=row[19],
+                release_date=row[20],
+                training_data_cutoff=row[21],
+                deprecation_date=row[22],
+                metadata=row[23],
+                last_scraped_at=row[24],
+                is_active=row[25],
+                is_deprecated=row[26],
+                created_at=row[27],
+                updated_at=row[28],
+            )
+
+    async def update_model_deprecation(self, arg: UpdateModelDeprecationParams) -> None:
+        await self._conn.execute(
+            sqlalchemy.text(UPDATE_MODEL_DEPRECATION),
+            {
+                "p1": arg.deprecation_date,
+                "p2": arg.replacement_model_id,
+                "p3": arg.slug,
+            },
+        )
+
+    async def upsert_model(self, arg: UpsertModelParams) -> Optional[models.LlmModel]:
+        row = (
+            await self._conn.execute(
+                sqlalchemy.text(UPSERT_MODEL),
+                {
+                    "p1": arg.provider_id,
+                    "p2": arg.model_id,
+                    "p3": arg.slug,
+                    "p4": arg.name,
+                    "p5": arg.display_name,
+                    "p6": arg.description,
+                    "p7": arg.version,
+                    "p8": arg.classification,
+                    "p9": arg.architecture,
+                    "p10": arg.capabilities,
+                    "p11": arg.performance,
+                    "p12": arg.token_info,
+                    "p13": arg.pricing,
+                    "p14": arg.api_details,
+                    "p15": arg.training,
+                    "p16": arg.licensing,
+                    "p17": arg.safety,
+                    "p18": arg.availability,
+                    "p19": arg.technical_specs,
+                    "p20": arg.release_date,
+                    "p21": arg.training_data_cutoff,
+                    "p22": arg.deprecation_date,
+                    "p23": arg.metadata,
+                },
+            )
+        ).first()
+        if row is None:
+            return None
+        return models.LlmModel(
+            id=row[0],
+            provider_id=row[1],
+            model_id=row[2],
+            slug=row[3],
+            name=row[4],
+            display_name=row[5],
+            description=row[6],
+            version=row[7],
+            capabilities=row[8],
+            classification=row[9],
+            architecture=row[10],
+            performance=row[11],
+            token_info=row[12],
+            pricing=row[13],
+            api_details=row[14],
+            training=row[15],
+            licensing=row[16],
+            safety=row[17],
+            availability=row[18],
+            technical_specs=row[19],
+            release_date=row[20],
+            training_data_cutoff=row[21],
+            deprecation_date=row[22],
+            metadata=row[23],
+            last_scraped_at=row[24],
+            is_active=row[25],
+            is_deprecated=row[26],
+            created_at=row[27],
+            updated_at=row[28],
+        )