feat(py): Core VertexAI plugin

firebase · Feb 19, 2025 · d13a10d · d13a10d
1 parent 22f937f
commit d13a10d
Show file tree

Hide file tree

Showing 9 changed files with 802 additions and 89 deletions.
diff --git a/captainhook.json b/captainhook.json
@@ -33,9 +33,6 @@
         {
           "run": "pnpm i --frozen-lockfile"
         },
-        {
-          "run": "py/bin/generate_schema_types"
-        },
         {
           "run": "bin/fmt"
         },
@@ -73,9 +70,6 @@
         {
           "run": "pnpm i --frozen-lockfile"
         },
-        {
-          "run": "py/bin/generate_schema_types"
-        },
         {
           "run": "bin/fmt"
         },

diff --git a/py/packages/genkit/src/genkit/ai/model.py b/py/packages/genkit/src/genkit/ai/model.py
@@ -3,6 +3,17 @@
 
 from collections.abc import Callable
 
-from genkit.core.schema_types import GenerateRequest, GenerateResponse
+from genkit.core.schemas import (
+    GenerateRequest,
+    GenerateResponse,
+    ModelInfo,
+)
+from pydantic.dataclasses import dataclass
 
 ModelFn = Callable[[GenerateRequest], GenerateResponse]
+
+
+@dataclass
+class ModelReference:
+    name: str
+    info: ModelInfo