From 0169dc54b0d4484ed5cda6031d8dbc05e847f6d6 Mon Sep 17 00:00:00 2001 From: Jeff Huang Date: Thu, 29 Jan 2026 16:59:19 -0600 Subject: [PATCH 1/4] refactor(py): move session/chat into blocks and remove public API --- py/packages/genkit/src/genkit/__init__.py | 6 - py/packages/genkit/src/genkit/ai/__init__.py | 5 - py/packages/genkit/src/genkit/ai/_aio.py | 298 +------------- py/packages/genkit/src/genkit/ai/_registry.py | 5 +- .../genkit/src/genkit/blocks/interfaces.py | 107 +++++ .../genkit/src/genkit/blocks/prompt.py | 8 +- .../src/genkit/blocks/session/__init__.py} | 9 +- .../src/genkit/{ => blocks}/session/chat.py | 120 ++---- .../genkit/{ => blocks}/session/in_memory.py | 0 .../genkit/{ => blocks}/session/session.py | 86 ++-- .../src/genkit/{ => blocks}/session/store.py | 0 .../genkit/src/genkit/blocks/session/types.py | 42 ++ .../genkit/src/genkit/session/__init__.py | 33 -- py/samples/chat-demo/README.md | 81 ---- py/samples/chat-demo/pyproject.toml | 69 ---- py/samples/chat-demo/run.sh | 24 -- py/samples/chat-demo/src/main.py | 355 ----------------- py/samples/session-demo/README.md | 51 --- py/samples/session-demo/pyproject.toml | 66 ---- py/samples/session-demo/src/main.py | 372 ------------------ 20 files changed, 233 insertions(+), 1504 deletions(-) create mode 100644 py/packages/genkit/src/genkit/blocks/interfaces.py rename py/{samples/session-demo/run.sh => packages/genkit/src/genkit/blocks/session/__init__.py} (75%) mode change 100755 => 100644 rename py/packages/genkit/src/genkit/{ => blocks}/session/chat.py (79%) rename py/packages/genkit/src/genkit/{ => blocks}/session/in_memory.py (100%) rename py/packages/genkit/src/genkit/{ => blocks}/session/session.py (83%) rename py/packages/genkit/src/genkit/{ => blocks}/session/store.py (100%) create mode 100644 py/packages/genkit/src/genkit/blocks/session/types.py delete mode 100644 py/packages/genkit/src/genkit/session/__init__.py delete mode 100644 py/samples/chat-demo/README.md delete mode 100644 py/samples/chat-demo/pyproject.toml delete mode 100755 py/samples/chat-demo/run.sh delete mode 100644 py/samples/chat-demo/src/main.py delete mode 100644 py/samples/session-demo/README.md delete mode 100644 py/samples/session-demo/pyproject.toml delete mode 100644 py/samples/session-demo/src/main.py diff --git a/py/packages/genkit/src/genkit/__init__.py b/py/packages/genkit/src/genkit/__init__.py index 334d206388..e01c123844 100644 --- a/py/packages/genkit/src/genkit/__init__.py +++ b/py/packages/genkit/src/genkit/__init__.py @@ -38,9 +38,6 @@ async def hello(name: str) -> str: GENKIT_VERSION, ActionKind, ActionRunContext, - Chat, - ChatOptions, - ChatStreamResponse, ExecutablePrompt, FlowWrapper, GenerateStreamResponse, @@ -101,9 +98,6 @@ async def hello(name: str) -> str: # From genkit.ai 'ActionKind', 'ActionRunContext', - 'Chat', - 'ChatOptions', - 'ChatStreamResponse', 'ExecutablePrompt', 'FlowWrapper', 'GenerateStreamResponse', diff --git a/py/packages/genkit/src/genkit/ai/__init__.py b/py/packages/genkit/src/genkit/ai/__init__.py index e8800bcd43..2570662388 100644 --- a/py/packages/genkit/src/genkit/ai/__init__.py +++ b/py/packages/genkit/src/genkit/ai/__init__.py @@ -33,7 +33,6 @@ from genkit.core.action import ActionRunContext from genkit.core.action.types import ActionKind from genkit.core.plugin import Plugin -from genkit.session import Chat, ChatOptions, ChatStreamResponse from ._aio import Genkit, Input, Output from ._registry import FlowWrapper, GenkitRegistry, SimpleRetrieverOptions @@ -61,10 +60,6 @@ 'OutputOptions', 'PromptGenerateOptions', 'ResumeOptions', - # Session/Chat - 'Chat', - 'ChatOptions', - 'ChatStreamResponse', # Document 'Document', # Plugin diff --git a/py/packages/genkit/src/genkit/ai/_aio.py b/py/packages/genkit/src/genkit/ai/_aio.py index a0cb4c773d..62f9fb1492 100644 --- a/py/packages/genkit/src/genkit/ai/_aio.py +++ b/py/packages/genkit/src/genkit/ai/_aio.py @@ -33,10 +33,7 @@ import uuid from collections.abc import AsyncIterator, Awaitable, Callable from pathlib import Path -from typing import TYPE_CHECKING, Any, Generic, TypedDict, TypeVar, cast, overload # noqa: F401 - -if TYPE_CHECKING: - from genkit.blocks.prompt import ExecutablePrompt +from typing import Any, TypeVar, cast, overload # noqa: F401 from opentelemetry import trace as trace_api from opentelemetry.sdk.trace import TracerProvider @@ -56,6 +53,8 @@ ModelMiddleware, ) from genkit.blocks.prompt import PromptConfig, load_prompt_folder, to_generate_action_options +from genkit.blocks.interfaces import Input as _Input +from genkit.blocks.interfaces import Output, OutputConfigDict from genkit.blocks.retriever import IndexerRef, IndexerRequest, RetrieverRef from genkit.core.action import Action, ActionRunContext from genkit.core.action.types import ActionKind @@ -70,7 +69,6 @@ Operation, SpanMetadata, ) -from genkit.session import Chat, ChatOptions, InMemorySessionStore, Session, SessionStore from genkit.types import ( DocumentData, GenerationCommonConfig, @@ -86,135 +84,13 @@ from ._server import ServerSpec T = TypeVar('T') - - -class OutputConfigDict(TypedDict, total=False): - """TypedDict for output configuration when passed as a dict.""" - - format: str | None - content_type: str | None - instructions: bool | str | None - schema: type | dict[str, object] | None - constrained: bool | None +Input = _Input InputT = TypeVar('InputT') OutputT = TypeVar('OutputT') -class Input(Generic[InputT]): - """Typed input configuration that preserves schema type information. - - This class provides a type-safe way to configure input schemas for prompts. - When you pass a Pydantic model as the schema, the prompt's input parameter - will be properly typed. - - Example: - ```python - from pydantic import BaseModel - - - class RecipeInput(BaseModel): - dish: str - servings: int - - - class Recipe(BaseModel): - name: str - ingredients: list[str] - - - # With Input[T] and Output[T], both input and output are typed - recipe_prompt = ai.define_prompt( - name='recipe', - prompt='Create a recipe for {dish} serving {servings} people', - input=Input(schema=RecipeInput), - output=Output(schema=Recipe), - ) - - # Input is type-checked! - response = await recipe_prompt(RecipeInput(dish='pizza', servings=4)) - response.output.name # ✓ Type checker knows this is str - ``` - - Attributes: - schema: The type/class for the input (Pydantic model, dataclass, etc.) - """ - - def __init__(self, schema: type[InputT]) -> None: - """Initialize typed input configuration. - - Args: - schema: The type/class for structured input. - """ - self.schema: type[InputT] = schema - - -class Output(Generic[OutputT]): - """Typed output configuration that preserves schema type information. - - This class provides a type-safe way to configure output options for generate(). - When you pass a Pydantic model or other type as the schema, the return type - of generate() will be properly typed. - - Example: - ```python - from pydantic import BaseModel - - - class Recipe(BaseModel): - name: str - ingredients: list[str] - - - # With Output[T], response.output is typed as Recipe - response = await ai.generate(prompt='Give me a pasta recipe', output=Output(schema=Recipe, format='json')) - response.output.name # ✓ Type checker knows this is str - ``` - - Attributes: - schema: The type/class for the output (Pydantic model, dataclass, etc.) - format: Output format name (e.g., 'json', 'text'). Defaults to 'json'. - content_type: MIME content type for the output. - instructions: Formatting instructions (True for default, False to disable, or custom str). - constrained: Whether to constrain model output to the schema. - """ - - def __init__( - self, - schema: type[OutputT], - format: str = 'json', - content_type: str | None = None, - instructions: bool | str | None = None, - constrained: bool | None = None, - ) -> None: - """Initialize typed output configuration. - - Args: - schema: The type/class for structured output. - format: Output format name. Defaults to 'json'. - content_type: Optional MIME content type. - instructions: Optional formatting instructions. - constrained: Whether to constrain output to schema. - """ - self.schema: type[OutputT] = schema - self.format: str = format - self.content_type: str | None = content_type - self.instructions: bool | str | None = instructions - self.constrained: bool | None = constrained - - def to_dict(self) -> OutputConfigDict: - """Convert to OutputConfigDict for internal use.""" - result: OutputConfigDict = {'schema': self.schema, 'format': self.format} - if self.content_type is not None: - result['content_type'] = self.content_type - if self.instructions is not None: - result['instructions'] = self.instructions - if self.constrained is not None: - result['constrained'] = self.constrained - return result - - class Genkit(GenkitBase): """Genkit asyncio user-facing API.""" @@ -265,172 +141,6 @@ def _resolve_embedder_name(self, embedder: str | EmbedderRef | None) -> str: else: raise ValueError('Embedder must be specified as a string name or an EmbedderRef.') - def create_session( - self, - store: SessionStore | None = None, - initial_state: dict[str, Any] | None = None, - ) -> Session: - """Creates a new session for multi-turn conversations. - - **Overview:** - - Initializes a new `Session` instance, which manages conversation history - and state. By default, it uses an ephemeral `InMemorySessionStore`, but - you should provide a persistent store (e.g. Firestore, Redis) for - production use. - - **Args:** - store: The `SessionStore` implementation to use. Defaults to `InMemorySessionStore`. - initial_state: A dictionary of initial state to populate the session with. - - **Returns:** - A new `Session` object bound to this Genkit instance. - - **Examples:** - - ```python - # ephemeral session - session = ai.create_session() - - # persistent session with initial state - session = ai.create_session(store=my_firestore_store, initial_state={'username': 'jdoe'}) - await session.chat('Hello') - ``` - """ - if store is None: - store = InMemorySessionStore() - - # pyrefly: ignore[bad-argument-type] - Self type is compatible with Genkit - session = Session(ai=self, store=store) - if initial_state: - session.update_state(initial_state) - return session - - async def load_session( - self, - session_id: str, - store: SessionStore, - ) -> Session | None: - """Loads an existing session from a store. - - **Overview:** - - Retrieves session data (history and state) from the provided `SessionStore` - and reconstructs a `Session` object. If the session ID is not found, - returns `None`. - - **Args:** - session_id: The unique identifier of the session to load. - store: The `SessionStore` to query. - - **Returns:** - The loaded `Session` object, or `None` if not found. - - **Examples:** - - ```python - session = await ai.load_session('sess_12345', store=my_store) - if session: - await session.chat('Continue our conversation') - else: - print('Session not found') - ``` - """ - data = await store.get(session_id) - if not data: - return None - # pyrefly: ignore[bad-argument-type] - Self type is compatible with Genkit - return Session(ai=self, store=store, data=data) - - def chat( - self, - preamble_or_options: ExecutablePrompt | ChatOptions | None = None, - options: ChatOptions | None = None, - ) -> Chat: - r"""Creates a chat session for multi-turn conversations (matches JS API). - - This method creates a Session and returns a Chat object for - conversational AI. It matches the JavaScript `ai.chat()` API exactly. - - Args: - preamble_or_options: Either an ExecutablePrompt to use as the - conversation preamble, or a ChatOptions dict with system - prompt, model, config, etc. - options: Additional ChatOptions (only used when first arg is - an ExecutablePrompt). - - Returns: - A Chat instance ready for multi-turn conversation. - - Example: - Basic chat with system prompt: - - ```python - chat = ai.chat({'system': 'You are a helpful pirate.'}) - response = await chat.send('Hello!') - print(response.text) # "Ahoy, matey!" - ``` - - Using an ExecutablePrompt: - - ```python - support_agent = ai.define_prompt( - name='support', - system='You are a customer support agent.', - ) - chat = ai.chat(support_agent) - response = await chat.send('My order is late') - ``` - - Streaming responses: - - ```python - chat = ai.chat({'system': 'Be verbose.'}) - result = chat.send_stream('Explain quantum physics') - - async for chunk in result.stream: - print(chunk.text, end='', flush=True) - - final = await result.response - ``` - - Multiple threads (use session.chat for thread names): - - ```python - session = ai.create_session() - lawyer = session.chat('lawyer', {'system': 'Talk like a lawyer'}) - pirate = session.chat('pirate', {'system': 'Talk like a pirate'}) - await lawyer.send('Tell me a joke') - await pirate.send('Tell me a joke') - ``` - - See Also: - - session.chat(): For named threads within a session - - JavaScript ai.chat(): js/genkit/src/genkit-beta.ts - """ - from genkit.blocks.prompt import ExecutablePrompt - - # Resolve preamble and options (matching JS pattern exactly) - preamble: ExecutablePrompt | None = None - chat_options: ChatOptions | None = None - - if preamble_or_options is not None: - if isinstance(preamble_or_options, ExecutablePrompt): - preamble = preamble_or_options - chat_options = options - else: - chat_options = preamble_or_options - - # Create session and use session.chat() (matches JS) - store = chat_options.get('store') if chat_options else None - session = self.create_session(store=store) - - if preamble is not None: - return session.chat(preamble, chat_options) - elif chat_options: - return session.chat(chat_options) - else: - return session.chat() @overload async def generate( diff --git a/py/packages/genkit/src/genkit/ai/_registry.py b/py/packages/genkit/src/genkit/ai/_registry.py index 078b8f6d9e..a8200d6cb1 100644 --- a/py/packages/genkit/src/genkit/ai/_registry.py +++ b/py/packages/genkit/src/genkit/ai/_registry.py @@ -43,15 +43,14 @@ import uuid from collections.abc import AsyncIterator, Awaitable, Callable from functools import wraps -from typing import TYPE_CHECKING, Any, Generic, ParamSpec, cast, overload +from typing import Any, Generic, ParamSpec, cast, overload from pydantic import BaseModel from typing_extensions import Never, TypeVar from genkit.aio import ensure_async -if TYPE_CHECKING: - from genkit.ai._aio import Input, Output +from genkit.blocks.interfaces import Input, Output from genkit.blocks.prompt import ExecutablePrompt from genkit.blocks.resource import FlexibleResourceFn, ResourceOptions diff --git a/py/packages/genkit/src/genkit/blocks/interfaces.py b/py/packages/genkit/src/genkit/blocks/interfaces.py new file mode 100644 index 0000000000..54c28cc587 --- /dev/null +++ b/py/packages/genkit/src/genkit/blocks/interfaces.py @@ -0,0 +1,107 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Shared interfaces and typing helpers across blocks.""" + +from __future__ import annotations + +from typing import Any, Generic, Protocol, TypedDict, TypeVar + +InputT = TypeVar('InputT') +OutputT = TypeVar('OutputT') + + +class OutputConfigDict(TypedDict, total=False): + """TypedDict for output configuration when passed as a dict.""" + + format: str | None + content_type: str | None + instructions: bool | str | None + schema: type | dict[str, object] | None + constrained: bool | None + + +class Input(Generic[InputT]): + """Typed input configuration that preserves schema type information. + + This class provides a type-safe way to configure input schemas for prompts. + When you pass a Pydantic model as the schema, the prompt's input parameter + will be properly typed. + """ + + def __init__(self, schema: type[InputT]) -> None: + """Initialize typed input configuration. + + Args: + schema: The type/class for structured input. + """ + + self.schema: type[InputT] = schema + + +class Output(Generic[OutputT]): + """Typed output configuration that preserves schema type information. + + This class provides a type-safe way to configure output options for generate(). + When you pass a Pydantic model or other type as the schema, the return type + of generate() will be properly typed. + """ + + def __init__( + self, + schema: type[OutputT], + format: str = 'json', + content_type: str | None = None, + instructions: bool | str | None = None, + constrained: bool | None = None, + ) -> None: + """Initialize typed output configuration. + + Args: + schema: The type/class for structured output. + format: Output format name. Defaults to 'json'. + content_type: Optional MIME content type. + instructions: Optional formatting instructions. + constrained: Whether to constrain output to schema. + """ + + self.schema: type[OutputT] = schema + self.format: str = format + self.content_type: str | None = content_type + self.instructions: bool | str | None = instructions + self.constrained: bool | None = constrained + + def to_dict(self) -> OutputConfigDict: + """Convert to OutputConfigDict for internal use.""" + + result: OutputConfigDict = {'schema': self.schema, 'format': self.format} + if self.content_type is not None: + result['content_type'] = self.content_type + if self.instructions is not None: + result['instructions'] = self.instructions + if self.constrained is not None: + result['constrained'] = self.constrained + return result + + +class ExecutablePromptLike(Protocol): + """Minimal interface for prompt-like objects used in typing.""" + + async def __call__(self, *args: Any, **kwargs: Any) -> Any: ... + + def stream(self, *args: Any, **kwargs: Any) -> Any: ... + + async def render(self, *args: Any, **kwargs: Any) -> Any: ... diff --git a/py/packages/genkit/src/genkit/blocks/prompt.py b/py/packages/genkit/src/genkit/blocks/prompt.py index b73a7290b2..3eba11f5b3 100644 --- a/py/packages/genkit/src/genkit/blocks/prompt.py +++ b/py/packages/genkit/src/genkit/blocks/prompt.py @@ -131,8 +131,8 @@ import weakref from collections.abc import AsyncIterable, Awaitable, Callable from pathlib import Path -from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypedDict, TypeVar, cast, overload - +from typing import Any, ClassVar, Generic, TypedDict, TypeVar, cast, overload +from genkit.blocks.interfaces import Input, Output from dotpromptz.typing import ( DataArgument, PromptFunction, @@ -1175,10 +1175,6 @@ async def as_tool(self) -> Action: return action -if TYPE_CHECKING: - from genkit.ai._aio import Input, Output - - # Overload 1: Both input and output typed -> ExecutablePrompt[InputT, OutputT] @overload def define_prompt( diff --git a/py/samples/session-demo/run.sh b/py/packages/genkit/src/genkit/blocks/session/__init__.py old mode 100755 new mode 100644 similarity index 75% rename from py/samples/session-demo/run.sh rename to py/packages/genkit/src/genkit/blocks/session/__init__.py index f4346468d1..5abfa1afdc --- a/py/samples/session-demo/run.sh +++ b/py/packages/genkit/src/genkit/blocks/session/__init__.py @@ -1,4 +1,3 @@ -#!/usr/bin/env bash # Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -15,10 +14,4 @@ # # SPDX-License-Identifier: Apache-2.0 -set -e - -# Run the script from the project root directory -cd "$(dirname "$0")" - -# Install dependencies and run the main script -uv run --managed-python streamlit run src/main.py "$@" +"""Internal session modules for Genkit blocks.""" diff --git a/py/packages/genkit/src/genkit/session/chat.py b/py/packages/genkit/src/genkit/blocks/session/chat.py similarity index 79% rename from py/packages/genkit/src/genkit/session/chat.py rename to py/packages/genkit/src/genkit/blocks/session/chat.py index b390fb16c5..412edc58e8 100644 --- a/py/packages/genkit/src/genkit/session/chat.py +++ b/py/packages/genkit/src/genkit/blocks/session/chat.py @@ -99,7 +99,7 @@ See Also: - JavaScript Chat: js/ai/src/chat.ts - - Session: genkit/session/session.py + - Session: genkit/blocks/session/session.py """ # pyright: reportImportCycles=false @@ -107,16 +107,14 @@ import asyncio from collections.abc import AsyncIterable -from typing import TYPE_CHECKING, Any, TypedDict +from typing import Any, TypedDict from genkit.aio import Channel from genkit.blocks.model import GenerateResponseChunkWrapper, GenerateResponseWrapper from genkit.types import Message, Part from .store import MAIN_THREAD - -if TYPE_CHECKING: - from .session import Session +from .types import SessionLike class ChatOptions(TypedDict, total=False): @@ -248,7 +246,7 @@ class Chat: def __init__( self, - session: Session, + session: SessionLike, request_base: dict[str, Any] | None = None, *, thread: str = MAIN_THREAD, @@ -267,7 +265,7 @@ def __init__( thread: Thread name for this conversation (default: 'main'). messages: Initial messages (from session thread or provided). """ - self._session: Session = session + self._session: SessionLike = session self._request_base: dict[str, Any] = request_base or {} self._thread_name: str = thread @@ -279,7 +277,7 @@ def __init__( self._messages = list(session.get_messages(self._thread_name)) @property - def session(self) -> Session: + def session(self) -> SessionLike: """The underlying Session object (readonly, matches JS).""" return self._session @@ -298,42 +296,19 @@ async def send( prompt: str | Part | list[Part], **kwargs: Any, # noqa: ANN401 - Forwarding generation arguments. ) -> GenerateResponseWrapper: - """Send a message and get the complete response. - - This method sends a user message, generates a response, updates the - message history, and persists the session state. - - Args: - prompt: The user message (string or parts). - **kwargs: Additional arguments passed to ai.generate(). - - Returns: - The model's response. + """Send a message and get the complete response (matches JS chat.send).""" + # Append user message to history + self._messages.append(Message(role='user', content=[Part.from_text(prompt)])) - Example: - ```python - response = await chat.send('What is the weather?') - print(response.text) + # Merge base options and call-specific options + gen_options = {**self._request_base, **kwargs} + gen_options['messages'] = self._messages - # With additional options - response = await chat.send('Explain briefly', config={'temperature': 0.5}) - ``` - """ - # Build generation options from request_base (pre-rendered options) - gen_options: dict[str, Any] = { - **self._request_base, - **kwargs, - 'messages': self._messages, - 'prompt': prompt, - } - - # Generate using session's ai instance + # Use session's AI instance to generate response response = await self._session._ai.generate(**gen_options) # pyright: ignore[reportPrivateUsage] - # Update message history - self._messages = response.messages - - # Persist to session + # Append model response to history + self._messages.append(Message(role='assistant', content=response.content)) self._session.update_messages(self._thread_name, self._messages) await self._session.save() @@ -344,55 +319,32 @@ def send_stream( prompt: str | Part | list[Part], **kwargs: Any, # noqa: ANN401 - Forwarding generation arguments. ) -> ChatStreamResponse: - r"""Send a message and stream the response. - - This method sends a user message and returns a streaming response - that can be iterated for chunks while also awaiting the final result. - - Args: - prompt: The user message (string or parts). - **kwargs: Additional arguments passed to ai.generate(). - - Returns: - A ChatStreamResponse with stream and response properties. - - Example: - ```python - result = chat.send_stream('Tell me a long story') + """Send a message and stream the response (matches JS chat.sendStream).""" + # Append user message to history + self._messages.append(Message(role='user', content=[Part.from_text(prompt)])) - async for chunk in result.stream: - print(chunk.text, end='', flush=True) - - final = await result.response - print(f'\\nDone! Used {final.usage} tokens') - ``` - """ - channel: Channel[GenerateResponseChunkWrapper, GenerateResponseWrapper[Any]] = Channel() + # Merge base options and call-specific options + gen_options = {**self._request_base, **kwargs} + gen_options['messages'] = self._messages - async def _do_send() -> GenerateResponseWrapper[Any]: - # Build generation options from request_base with streaming callback - gen_options: dict[str, Any] = { - **self._request_base, - **kwargs, - 'messages': self._messages, - 'prompt': prompt, - 'on_chunk': lambda chunk: channel.send(chunk), - } + # Use session's AI instance to generate response (streaming) + stream_result = self._session._ai.generate_stream(**gen_options) # pyright: ignore[reportPrivateUsage] - # Generate using session's ai instance - response = await self._session._ai.generate(**gen_options) # pyright: ignore[reportPrivateUsage] - - # Update message history - self._messages = response.messages - - # Persist to session + async def _run_stream() -> GenerateResponseWrapper: + response = await stream_result.response + # Append model response to history + self._messages.append(Message(role='assistant', content=response.content)) self._session.update_messages(self._thread_name, self._messages) await self._session.save() - return response - # Create task and use set_close_future to close the channel when done - response_task = asyncio.create_task(_do_send()) - channel.set_close_future(response_task) + response_future = asyncio.ensure_future(_run_stream()) + return ChatStreamResponse(stream_result.stream, response_future) + - return ChatStreamResponse(channel=channel, response_future=response_task) +def _merge_chat_options(base: dict[str, Any], updates: dict[str, Any]) -> dict[str, Any]: + """Merge chat options, filtering out chat-only keys.""" + merged = {**base, **updates} + for key in _CHAT_ONLY_KEYS: + merged.pop(key, None) + return merged diff --git a/py/packages/genkit/src/genkit/session/in_memory.py b/py/packages/genkit/src/genkit/blocks/session/in_memory.py similarity index 100% rename from py/packages/genkit/src/genkit/session/in_memory.py rename to py/packages/genkit/src/genkit/blocks/session/in_memory.py diff --git a/py/packages/genkit/src/genkit/session/session.py b/py/packages/genkit/src/genkit/blocks/session/session.py similarity index 83% rename from py/packages/genkit/src/genkit/session/session.py rename to py/packages/genkit/src/genkit/blocks/session/session.py index fc4adcb40e..61c917c4e7 100644 --- a/py/packages/genkit/src/genkit/session/session.py +++ b/py/packages/genkit/src/genkit/blocks/session/session.py @@ -37,16 +37,14 @@ import time import uuid -from typing import TYPE_CHECKING, Any +from typing import Any +from genkit.blocks.interfaces import ExecutablePromptLike from genkit.types import Message -if TYPE_CHECKING: - from genkit.ai import Genkit - from genkit.blocks.prompt import ExecutablePrompt - from .chat import Chat, ChatOptions from .store import MAIN_THREAD, SessionData, SessionStore +from .types import GenkitLike class Session: @@ -102,7 +100,7 @@ class Session: def __init__( self, - ai: Genkit, + ai: GenkitLike, store: SessionStore, id: str | None = None, data: SessionData | None = None, @@ -115,25 +113,24 @@ def __init__( id: The session ID. If not provided, a new UUID is generated. data: Existing session data (if loading). """ - self._ai: Genkit = ai + self._ai: GenkitLike = ai self._store: SessionStore = store if data: - data_id = data.get('id') - self._id = data_id if data_id is not None else (id or str(uuid.uuid4())) - self._state: dict[str, Any] = data.get('state') or {} - # Load threads, with backward compat for legacy 'messages' field - self._threads: dict[str, list[Message]] = data.get('threads') or {} - if not self._threads and data.get('messages'): - # Migrate legacy messages to main thread - self._threads[MAIN_THREAD] = data.get('messages') or [] + self._id: str = data.get('id', str(uuid.uuid4())) + self._state: dict[str, Any] = data.get('state', {}) or {} + threads = data.get('threads', {}) or {} + if not threads and data.get('messages'): + # Backwards compatibility: use legacy messages field as main thread + threads = {MAIN_THREAD: data.get('messages', [])} + self._threads: dict[str, list[Message]] = threads self._created_at: float = data.get('created_at') or time.time() - self._updated_at: float = data.get('updated_at') or time.time() + self._updated_at: float = data.get('updated_at') or self._created_at else: self._id = id or str(uuid.uuid4()) self._state = {} - self._threads = {} + self._threads = {MAIN_THREAD: []} self._created_at = time.time() - self._updated_at = time.time() + self._updated_at = self._created_at @property def id(self) -> str: @@ -142,41 +139,44 @@ def id(self) -> str: @property def state(self) -> dict[str, Any]: - """The session state.""" + """Session state data (mutable).""" return self._state @property def messages(self) -> list[Message]: - """The main thread message history (for backward compatibility).""" - return self.get_messages(MAIN_THREAD) - - @property - def threads(self) -> dict[str, list[Message]]: - """All thread message histories.""" - return self._threads + """Messages for the main thread (legacy compatibility).""" + return self._threads.get(MAIN_THREAD, []) - def get_messages(self, thread: str = MAIN_THREAD) -> list[Message]: - """Get messages for a specific thread. + def get_messages(self, thread_name: str = MAIN_THREAD) -> list[Message]: + """Get messages for a thread. Args: - thread: The thread name (defaults to MAIN_THREAD). + thread_name: Name of the thread. Returns: - List of messages for the thread, or empty list if thread doesn't exist. + List of messages for the thread (empty list if thread not found). """ - return self._threads.get(thread, []) + return self._threads.get(thread_name, []) - def update_messages(self, thread: str, messages: list[Message]) -> None: - """Update messages for a specific thread. + def update_messages(self, thread_name: str, messages: list[Message]) -> None: + """Update messages for a thread. Args: - thread: The thread name. - messages: The new message list for the thread. + thread_name: Name of the thread. + messages: New list of messages. """ - self._threads[thread] = messages + self._threads[thread_name] = messages + + def update_state(self, updates: dict[str, Any]) -> None: + """Update session state. + + Args: + updates: Dictionary of state updates to apply. + """ + self._state.update(updates) async def save(self) -> None: - """Persist the current session state to the store.""" + """Persist the session to the store.""" self._updated_at = time.time() data: SessionData = { 'id': self._id, @@ -188,18 +188,10 @@ async def save(self) -> None: } await self._store.save(self._id, data) - def update_state(self, updates: dict[str, Any]) -> None: - """Update session state. - - Args: - updates: Dictionary of state updates to apply. - """ - self._state.update(updates) - def chat( self, - thread_or_preamble_or_options: str | ExecutablePrompt | ChatOptions | None = None, - preamble_or_options: ExecutablePrompt | ChatOptions | None = None, + thread_or_preamble_or_options: str | ExecutablePromptLike | ChatOptions | None = None, + preamble_or_options: ExecutablePromptLike | ChatOptions | None = None, options: ChatOptions | None = None, ) -> Chat: """Create a Chat object for multi-turn conversation (matching JS API). diff --git a/py/packages/genkit/src/genkit/session/store.py b/py/packages/genkit/src/genkit/blocks/session/store.py similarity index 100% rename from py/packages/genkit/src/genkit/session/store.py rename to py/packages/genkit/src/genkit/blocks/session/store.py diff --git a/py/packages/genkit/src/genkit/blocks/session/types.py b/py/packages/genkit/src/genkit/blocks/session/types.py new file mode 100644 index 0000000000..b08ae2e793 --- /dev/null +++ b/py/packages/genkit/src/genkit/blocks/session/types.py @@ -0,0 +1,42 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Shared session typing contracts to avoid import cycles.""" + +from __future__ import annotations + +from typing import Any, Protocol + +from genkit.types import Message + + +class GenkitLike(Protocol): + """Minimal Genkit surface needed by session/chat.""" + + async def generate(self, *args: Any, **kwargs: Any) -> Any: ... + + +class SessionLike(Protocol): + """Minimal Session surface needed by Chat.""" + + id: str + _ai: GenkitLike + + def get_messages(self, thread_name: str) -> list[Message]: ... + + def update_messages(self, thread_name: str, messages: list[Message]) -> None: ... + + async def save(self) -> None: ... diff --git a/py/packages/genkit/src/genkit/session/__init__.py b/py/packages/genkit/src/genkit/session/__init__.py deleted file mode 100644 index cc8f107cd9..0000000000 --- a/py/packages/genkit/src/genkit/session/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Session management for Genkit.""" - -from .chat import Chat, ChatOptions, ChatStreamResponse -from .in_memory import InMemorySessionStore -from .session import Session -from .store import MAIN_THREAD, SessionData, SessionStore - -__all__ = [ - 'Chat', - 'ChatOptions', - 'ChatStreamResponse', - 'InMemorySessionStore', - 'MAIN_THREAD', - 'Session', - 'SessionData', - 'SessionStore', -] diff --git a/py/samples/chat-demo/README.md b/py/samples/chat-demo/README.md deleted file mode 100644 index a5bff3d39f..0000000000 --- a/py/samples/chat-demo/README.md +++ /dev/null @@ -1,81 +0,0 @@ -# Genkit Chat Demo - -This sample demonstrates how to build a full-featured web chat application using Genkit's session and chat APIs with Streamlit. - -## Features - -* **Multi-Model Support**: Switch between Google (Gemini), Anthropic (Claude), DeepSeek, xAI (Grok), OpenAI, and Ollama. -* **Multiple Conversations**: Create and manage parallel chat threads. -* **Persistent Sessions**: Conversation history maintained per thread. -* **Streaming**: Real-time response streaming. -* **Rich UI**: Built with Streamlit for a web-based experience. - -## Running - -```bash -./run.sh -``` - -This launches the Streamlit app. Configure API keys in the sidebar. - -## Key APIs Demonstrated - -| API | Description | Example | -|-----|-------------|---------| -| `ai.chat()` | Create a chat with thread support | `chat = ai.chat(thread='conv1')` | -| `chat.send()` | Send message, get complete response | `response = await chat.send('Hi')` | -| `chat.send_stream()` | Send message, stream response | `result = chat.send_stream('Hi')` | -| `chat.thread` | Get the thread name | `print(chat.thread)` | - -## Prerequisites - -* Python 3.12+ -* `uv` (recommended) - -## Setup - -Set your API keys as environment variables: - -```bash -export GEMINI_API_KEY=your-key # For Google Gemini -export ANTHROPIC_API_KEY=your-key # For Anthropic Claude -export DEEPSEEK_API_KEY=your-key # For DeepSeek -export XAI_API_KEY=your-key # For xAI Grok -export OPENAI_API_KEY=your-key # For OpenAI -``` - -For Vertex AI, also set: - -```bash -export VERTEX_AI_PROJECT_ID=your-project -export VERTEX_AI_LOCATION=us-central1 -``` - -For Ollama, ensure your Ollama server is running locally. - -## Testing This Demo - -1. **Run the demo**: - ```bash - ./run.sh - ``` - -2. **Open the Streamlit UI** at http://localhost:8501 - -3. **Test the following**: - - [ ] Select different models from the sidebar dropdown - - [ ] Send a message and verify response appears - - [ ] Test streaming (responses should appear incrementally) - - [ ] Create multiple conversations using "New Conversation" - - [ ] Switch between conversations and verify history persists - - [ ] Test with different providers (Google, Anthropic, OpenAI, etc.) - -4. **Expected behavior**: - - Messages appear in chat bubbles - - Streaming shows text appearing word-by-word - - Each conversation maintains separate history - - Model switching works mid-conversation - -## Related Samples - -* `session-demo/` - Lower-level session management examples diff --git a/py/samples/chat-demo/pyproject.toml b/py/samples/chat-demo/pyproject.toml deleted file mode 100644 index a0bdb9a787..0000000000 --- a/py/samples/chat-demo/pyproject.toml +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -[project] -authors = [{ name = "Google" }] -classifiers = [ - "Development Status :: 3 - Alpha", - "Environment :: Console", - "Environment :: Web Environment", - "Intended Audience :: Developers", - "Operating System :: OS Independent", - "Programming Language :: Python", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", - "Programming Language :: Python :: 3.14", - "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries", -] -dependencies = [ - "genkit", - "genkit-plugin-anthropic", - "genkit-plugin-compat-oai", - "genkit-plugin-deepseek", - "genkit-plugin-google-cloud", - "genkit-plugin-google-genai", - "genkit-plugin-ollama", - "genkit-plugin-vertex-ai", - "genkit-plugin-xai", - "ollama", - "pydantic>=2.10.5", - "streamlit>=1.40.0", - "structlog>=25.2.0", - "uvloop>=0.21.0", -] -description = "Chat demo using session management" -license = "Apache-2.0" -name = "chat-demo" -readme = "README.md" -requires-python = ">=3.10" -version = "0.1.0" - -[project.optional-dependencies] -dev = ["watchdog>=6.0.0"] - -[build-system] -build-backend = "hatchling.build" -requires = ["hatchling"] - -[tool.hatch.build.targets.wheel] -packages = ["src/main.py"] - -[tool.hatch.metadata] -allow-direct-references = true diff --git a/py/samples/chat-demo/run.sh b/py/samples/chat-demo/run.sh deleted file mode 100755 index f4346468d1..0000000000 --- a/py/samples/chat-demo/run.sh +++ /dev/null @@ -1,24 +0,0 @@ -#!/usr/bin/env bash -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -set -e - -# Run the script from the project root directory -cd "$(dirname "$0")" - -# Install dependencies and run the main script -uv run --managed-python streamlit run src/main.py "$@" diff --git a/py/samples/chat-demo/src/main.py b/py/samples/chat-demo/src/main.py deleted file mode 100644 index 24d0d8696f..0000000000 --- a/py/samples/chat-demo/src/main.py +++ /dev/null @@ -1,355 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Chat Demo with Streamlit UI. - -This sample implements a web-based chatbot using Streamlit. It uses the -ai.chat() convenience API with threads for managing multiple conversations. - -See README.md for testing instructions. - -Key Features -============ -| Feature Description | Example Usage | -|-----------------------------------------|-----------------------------------------| -| ai.chat() convenience API | `chat = ai.chat(system='...')` | -| Chat.send() for messages | `response = await chat.send('Hi')` | -| Chat.send_stream() for streaming | `result = chat.send_stream('Hi')` | -| Thread-based conversations | `ai.chat(thread='conv1')` | -| Multiple parallel conversations | Each conversation is a separate thread | -""" - -import asyncio -import logging -import os - -import streamlit as st - -from genkit.ai import Genkit -from genkit.blocks.model import GenerateResponseChunkWrapper, GenerateResponseWrapper, Message -from genkit.core.logging import get_logger -from genkit.core.typing import Part, TextPart -from genkit.plugins.anthropic import Anthropic -from genkit.plugins.compat_oai import OpenAI -from genkit.plugins.deepseek import DeepSeek -from genkit.plugins.google_genai import GoogleAI, VertexAI -from genkit.plugins.ollama import Ollama -from genkit.plugins.vertex_ai import ModelGardenPlugin -from genkit.plugins.xai import XAI - -logging.basicConfig(level=logging.INFO) -logger = get_logger(__name__) - -st.set_page_config(page_title='Genkit Chat Demo', page_icon='💬') -st.title('💬 Genkit Chat Demo') - -# --- Model Definitions --- -DEFAULT_MODEL = 'googleai/gemini-3-flash-preview' - -AVAILABLE_MODELS = { - 'Google (Gemini API)': [ - 'googleai/gemini-3-flash-preview', - 'googleai/gemini-3-pro-preview', - 'googleai/gemini-2.5-pro', - 'googleai/gemini-2.0-flash', - 'googleai/gemini-2.0-pro-exp-02-05', - ], - 'Vertex AI (Gemini)': [ - 'vertexai/gemini-2.5-pro', - 'vertexai/gemini-2.5-flash', - 'vertexai/gemini-2.0-flash', - 'vertexai/gemini-2.0-flash-lite', - ], - 'Anthropic': [ - 'anthropic/claude-sonnet-4-5', - 'anthropic/claude-opus-4-5', - 'anthropic/claude-haiku-4-5', - 'anthropic/claude-3-5-sonnet', - ], - 'DeepSeek': [ - 'deepseek/deepseek-chat', - 'deepseek/deepseek-reasoner', - ], - 'xAI': [ - 'xai/grok-4', - 'xai/grok-3', - 'xai/grok-2-latest', - 'xai/grok-2-vision-1212', - 'xai/grok-2-1212', - 'xai/grok-beta', - ], - 'OpenAI': [ - 'openai/gpt-5.1', - 'openai/gpt-5', - 'openai/gpt-4o', - 'openai/gpt-4o-mini', - 'openai/o3', - 'openai/o1', - ], - 'Vertex AI Model Garden': [ - 'modelgarden/anthropic/claude-3-5-sonnet-v2@20241022', - 'modelgarden/anthropic/claude-3-opus@20240229', - 'modelgarden/anthropic/claude-3-sonnet@20240229', - 'modelgarden/anthropic/claude-3-haiku@20240307', - 'modelgarden/meta/llama-3.1-405b-instruct-maas', - 'modelgarden/meta/llama-3.2-90b-vision-instruct-maas', - ], - 'Ollama': [], # Dynamically populated -} - -# Fetch Ollama models if provider selected or just list them if server is running -try: - import ollama - - try: - ollama_client = ollama.Client() - models_resp = ollama_client.list() - AVAILABLE_MODELS['Ollama'] = [f'ollama/{m.model}' for m in models_resp.models] - except Exception: - AVAILABLE_MODELS['Ollama'] = ['ollama/llama3.2', 'ollama/gemma2'] -except ImportError: - AVAILABLE_MODELS['Ollama'] = ['ollama/llama3.2', 'ollama/gemma2'] - -# Flatten available models list for default selection -all_models = [m for models in AVAILABLE_MODELS.values() for m in models] -default_idx = all_models.index(DEFAULT_MODEL) if DEFAULT_MODEL in all_models else 0 - -# --- State Initialization --- -if 'selected_provider' not in st.session_state: - st.session_state['selected_provider'] = list(AVAILABLE_MODELS.keys())[0] -if 'selected_model' not in st.session_state: - st.session_state['selected_model'] = AVAILABLE_MODELS[st.session_state['selected_provider']][0] - - -def update_provider() -> None: - """Reset model when provider changes to avoids errors.""" - pass - - -# --- Sidebar Configuration --- -api_keys = {} -PROVIDERS = { - 'Google (Gemini API)': 'GEMINI_API_KEY', - 'Anthropic': 'ANTHROPIC_API_KEY', - 'DeepSeek': 'DEEPSEEK_API_KEY', - 'xAI': 'XAI_API_KEY', - 'OpenAI': 'OPENAI_API_KEY', - 'Vertex AI Project': 'VERTEX_AI_PROJECT_ID', - 'Vertex AI Location': 'VERTEX_AI_LOCATION', -} - -with st.sidebar: - # 1. Model Selection (Pinned to Top) - st.header('Model Selection') - provider = st.selectbox( - 'Provider', - list(AVAILABLE_MODELS.keys()), - key='selected_provider', - on_change=update_provider, - ) - - curr_prov = st.session_state['selected_provider'] - # Ensure selected model is valid for provider - valid_models = AVAILABLE_MODELS[curr_prov] - if st.session_state['selected_model'] not in valid_models: - st.session_state['selected_model'] = valid_models[0] - - st.selectbox( - 'Model', - valid_models, - key='selected_model', - ) - st.caption(f'Selected: `{st.session_state["selected_model"]}`') - - enable_streaming = st.checkbox('Enable Streaming', value=True) - st.divider() - - # 2. Authentication - st.header('Authentication') - for label, env_var in PROVIDERS.items(): - val = os.environ.get(env_var) - if not val: - try: - val = st.secrets.get(env_var) - except Exception: - val = None - - if val: - st.success(f'{label}: Configured', icon='✅') - api_keys[env_var] = val - else: - user_key = st.text_input(f'{label} API Key', type='password', help=f'Set {env_var}') - if user_key: - os.environ[env_var] = user_key - api_keys[env_var] = user_key - st.rerun() - else: - st.warning(f'{label}: Not set', icon='⚠️') - api_keys[env_var] = None - st.divider() - -# --- Main Logic --- - -# Init Genkit with current state values -provider_val = st.session_state['selected_provider'] -model_val = st.session_state['selected_model'] - -# Initialize plugins based on available keys -plugins = [] -if api_keys.get('GEMINI_API_KEY'): - plugins.append(GoogleAI()) -if api_keys.get('ANTHROPIC_API_KEY'): - plugins.append(Anthropic()) -if api_keys.get('DEEPSEEK_API_KEY'): - plugins.append(DeepSeek()) -if api_keys.get('XAI_API_KEY'): - plugins.append(XAI()) -if api_keys.get('OPENAI_API_KEY'): - plugins.append(OpenAI()) - -# Always add Ollama if it has models or just by default (no auth needed typically) -if provider_val == 'Ollama': - plugins.append(Ollama()) - -if api_keys.get('VERTEX_AI_PROJECT_ID') and api_keys.get('VERTEX_AI_LOCATION'): - # Add Model Garden Support - plugins.append( - ModelGardenPlugin(project_id=api_keys['VERTEX_AI_PROJECT_ID'], location=api_keys['VERTEX_AI_LOCATION']) - ) - # Add Vertex AI Gemini Support - plugins.append(VertexAI(project=api_keys['VERTEX_AI_PROJECT_ID'], location=api_keys['VERTEX_AI_LOCATION'])) - -if not plugins and provider_val != 'Ollama': - st.error('No API keys provided for selected provider.') - st.stop() - -ai = Genkit( - plugins=plugins, - model=model_val, -) - -# Initialize conversation threads -# We use a simple dict to track thread names and their display messages -if 'threads' not in st.session_state: - st.session_state['threads'] = ['Conversation 1'] - st.session_state['active_thread'] = 'Conversation 1' - st.session_state['thread_messages'] = {'Conversation 1': []} # For UI display - -# 3. Conversations (Sidebar - Appended) -with st.sidebar: - st.header('Conversations') - - if st.button('➕ New Conversation', use_container_width=True): - new_thread = f'Conversation {len(st.session_state["threads"]) + 1}' - st.session_state['threads'].append(new_thread) - st.session_state['thread_messages'][new_thread] = [] - st.session_state['active_thread'] = new_thread - st.rerun() - - # Conversation List - thread_names = st.session_state['threads'] - - # Determine index dynamically to handle deletions or state changes safely - try: - idx = thread_names.index(st.session_state['active_thread']) - except ValueError: - idx = 0 - - selected_thread = st.radio( - 'History', - options=thread_names, - index=idx, - label_visibility='collapsed', - ) - - if selected_thread != st.session_state['active_thread']: - st.session_state['active_thread'] = selected_thread - st.rerun() - - if st.button('🗑️ Clear Current History', use_container_width=True): - st.session_state['thread_messages'][st.session_state['active_thread']] = [] - st.rerun() - - -# Get the current thread name and messages for display -current_thread = st.session_state['active_thread'] -messages = st.session_state['thread_messages'][current_thread] - -# Display chat messages from history -for message in messages: - with st.chat_message(message['role']): - st.markdown(message['content']) - -# Accept user input -if prompt := st.chat_input('What is up?'): - # Add user message to chat history UI - messages.append({'role': 'user', 'content': prompt}) - with st.chat_message('user'): - st.markdown(prompt) - - # Generate response - with st.chat_message('assistant'): - message_placeholder = st.empty() - full_response = '' - - async def run_chat() -> GenerateResponseWrapper: - """Run chat using ai.chat() with thread support.""" - # Build message history from stored conversation (excluding the prompt we just added) - history: list[Message] = [] - for msg_dict in messages[:-1]: # Exclude last message (the prompt we just added) - role = 'model' if msg_dict['role'] == 'assistant' else msg_dict['role'] - history.append(Message(role=role, content=[Part(root=TextPart(text=msg_dict['content']))])) - - # Create a chat with the history restored - # Note: Using ChatOptions dict to pass messages (matches JS API) - chat = ai.chat({'messages': history}) - - # Streaming callback - full_text = '' - - def on_chunk(chunk: GenerateResponseChunkWrapper) -> None: - nonlocal full_text - if hasattr(chunk, 'text') and chunk.text: - full_text += chunk.text - message_placeholder.markdown(full_text + '▌') - - # Send the message - assert prompt is not None # Guaranteed by walrus operator guard - if enable_streaming: - result = chat.send_stream(prompt) - async for chunk in result.stream: - on_chunk(chunk) - return await result.response - else: - return await chat.send(prompt) - - # Use asyncio.run which handles the loop correctly for this thread - try: - response = asyncio.run(run_chat()) - full_response = response.text - message_placeholder.markdown(full_response) - messages.append({'role': 'assistant', 'content': full_response}) - except Exception as e: - # Check if it's a GenkitError or other exceptions and extract nice message - error_msg = str(e) - if hasattr(e, 'message'): - error_msg = e.message - elif hasattr(e, 'cause') and e.cause: - error_msg = str(e.cause) - - st.error(f'Error: {error_msg}') - # Stop spinner/execution - st.stop() diff --git a/py/samples/session-demo/README.md b/py/samples/session-demo/README.md deleted file mode 100644 index 25750f02f3..0000000000 --- a/py/samples/session-demo/README.md +++ /dev/null @@ -1,51 +0,0 @@ -# Genkit Session Management Demo - -This sample demonstrates how to use Genkit's advanced session management features to create stateful conversations with structured state updates. - -## Features Demonstrated - -* **Interactive Flow**: A guided step-by-step interactive flow using Streamlit. -* **Session Creation**: Initializing sessions with initial state (e.g., `user_name`, `preference`). -* **State Updates**: Updating session state dynamically during the conversation (e.g., changing `topic`). -* **Multi-Provider Support**: Switch between models from Google, Anthropic, DeepSeek, xAI, OpenAI, and Ollama. -* **State Visualization**: Real-time view of the session's internal stateJSON. - -## Prerequisites - -* Python 3.12+ -* `uv` (recommended) - -## Setup - -1. (Optional) Set your API keys as environment variables for convenience: - * `GEMINI_API_KEY` - * `ANTHROPIC_API_KEY` - * `DEEPSEEK_API_KEY` - * `XAI_API_KEY` - * `OPENAI_API_KEY` -2. If using Ollama, ensure your Ollama server is running locally. - -## Running - -```bash -./run.sh -``` - -This will launch the Streamlit app. Follow the interactive buttons to step through the session flow and watch the state update in real-time. - -## Testing This Demo - -1. **Open the Streamlit UI** at http://localhost:8501 - -2. **Test the following**: - - [ ] Create a new session and note the session ID - - [ ] Send messages and verify they're stored in session - - [ ] Check that session state is displayed/updated - - [ ] Reload the page and verify session persists (in-memory) - - [ ] Create multiple sessions and switch between them - - [ ] Test with different model providers - -3. **Expected behavior**: - - Sessions maintain conversation history - - State changes persist within session lifetime - - Each session has unique ID for identification diff --git a/py/samples/session-demo/pyproject.toml b/py/samples/session-demo/pyproject.toml deleted file mode 100644 index 4a0bcab80c..0000000000 --- a/py/samples/session-demo/pyproject.toml +++ /dev/null @@ -1,66 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -[project] -authors = [{ name = "Google" }] -classifiers = [ - "Development Status :: 3 - Alpha", - "Environment :: Console", - "Environment :: Web Environment", - "Intended Audience :: Developers", - "Operating System :: OS Independent", - "Programming Language :: Python", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", - "Programming Language :: Python :: 3.14", - "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries", -] -dependencies = [ - "genkit", - "genkit-plugin-anthropic", - "genkit-plugin-deepseek", - "genkit-plugin-google-genai", - "genkit-plugin-ollama", - "genkit-plugin-xai", - "ollama", - "pydantic>=2.10.5", - "streamlit>=1.40.0", - "structlog>=25.2.0", - "uvloop>=0.21.0", -] -description = "Session management demo" -license = "Apache-2.0" -name = "session-demo" -readme = "README.md" -requires-python = ">=3.10" -version = "0.1.0" - -[project.optional-dependencies] -dev = ["watchdog>=6.0.0"] - -[build-system] -build-backend = "hatchling.build" -requires = ["hatchling"] - -[tool.hatch.build.targets.wheel] -packages = ["src/main.py"] - -[tool.hatch.metadata] -allow-direct-references = true diff --git a/py/samples/session-demo/src/main.py b/py/samples/session-demo/src/main.py deleted file mode 100644 index ff09a25bc5..0000000000 --- a/py/samples/session-demo/src/main.py +++ /dev/null @@ -1,372 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Session Demo - Managing Conversation State with Genkit Sessions. - -This sample demonstrates how to use Genkit's session management capabilities -to maintain conversation state across multiple interactions. - -See README.md for testing instructions. - -Key Features -============ -| Feature | Description | -|----------------------------------|------------------------------------------| -| `InMemorySessionStore` | Store sessions in memory | -| `Session` object | Encapsulates conversation state | -| `session.create()` | Create new session with initial state | -| `session.load()` | Load existing session by ID | -| `session.state` | Access/modify session state dict | -| Multi-provider support | Works with Google, Anthropic, OpenAI | -""" - -import asyncio -import logging -import os -from typing import Any - -import streamlit as st - -from genkit.ai import Genkit -from genkit.blocks.model import GenerateResponseWrapper, Message -from genkit.core.logging import get_logger -from genkit.core.typing import Part, TextPart -from genkit.plugins.anthropic import Anthropic -from genkit.plugins.compat_oai import OpenAI -from genkit.plugins.deepseek import DeepSeek -from genkit.plugins.google_genai import GoogleAI -from genkit.plugins.ollama import Ollama -from genkit.plugins.xai import XAI -from genkit.session import InMemorySessionStore, Session - -logging.basicConfig(level=logging.INFO) -logger = get_logger(__name__) - -st.set_page_config(page_title='Genkit Session Demo', page_icon='📝', layout='wide') -st.title('📝 Genkit Session Demo') - -# Check API keys -PROVIDERS = { - 'Google': 'GEMINI_API_KEY', - 'Anthropic': 'ANTHROPIC_API_KEY', - 'DeepSeek': 'DEEPSEEK_API_KEY', - 'xAI': 'XAI_API_KEY', - 'OpenAI': 'OPENAI_API_KEY', -} - -api_keys = {} - -with st.sidebar: - st.header('Authentication') - for label, env_var in PROVIDERS.items(): - val = os.environ.get(env_var) - if not val: - try: - val = st.secrets.get(env_var) - except Exception: - val = None - - if val: - st.success(f'{label}: Configured', icon='✅') - api_keys[env_var] = val - else: - user_key = st.text_input(f'{label} API Key', type='password', help=f'Set {env_var}') - if user_key: - os.environ[env_var] = user_key - api_keys[env_var] = user_key - st.rerun() - else: - st.warning(f'{label}: Not set', icon='⚠️') - api_keys[env_var] = None - -# --- Initialization --- - -# Initialize simple session state variables if they don't exist -if 'step' not in st.session_state: - st.session_state['step'] = 0 - st.session_state['session_data'] = None - st.session_state['logs'] = [] # To show chat/logs - -# Default model -DEFAULT_MODEL = 'googleai/gemini-3-flash-preview' - -AVAILABLE_MODELS = { - 'Google': [ - 'googleai/gemini-3-flash-preview', - 'googleai/gemini-3-pro-preview', - 'googleai/gemini-2.5-pro', - 'googleai/gemini-2.0-flash', - 'googleai/gemini-2.0-pro-exp-02-05', - ], - 'Anthropic': [ - 'anthropic/claude-sonnet-4-5', - 'anthropic/claude-opus-4-5', - 'anthropic/claude-haiku-4-5', - 'anthropic/claude-3-5-sonnet', - ], - 'DeepSeek': [ - 'deepseek/deepseek-v3', - 'deepseek/deepseek-r1', - 'deepseek/deepseek-chat', - 'deepseek/deepseek-reasoner', - ], - 'xAI': [ - 'xai/grok-4.1', - 'xai/grok-4', - 'xai/grok-3', - 'xai/grok-2-latest', - 'xai/grok-2-vision-1212', - 'xai/grok-2-1212', - 'xai/grok-beta', - ], - 'OpenAI': [ - 'openai/gpt-4.5-preview', - 'openai/gpt-4o', - 'openai/gpt-4o-mini', - 'openai/o1-preview', - 'openai/o1-mini', - 'openai/o3-mini', - ], - 'Ollama': [], -} - -# Fetch Ollama models -try: - import ollama - - try: - ollama_client = ollama.Client() - models_resp = ollama_client.list() - AVAILABLE_MODELS['Ollama'] = [f'ollama/{m.model}' for m in models_resp.models] - except Exception: - AVAILABLE_MODELS['Ollama'] = ['ollama/llama3.2', 'ollama/gemma2'] -except ImportError: - AVAILABLE_MODELS['Ollama'] = ['ollama/llama3.2', 'ollama/gemma2'] - -# Sidebar for configuration -with st.sidebar: - st.header('Configuration') - - # Provider selection - provider = st.selectbox('Provider', list(AVAILABLE_MODELS.keys())) - - # Model selection based on provider - model_name = st.selectbox('Model', AVAILABLE_MODELS[provider], index=0) - st.caption(f'Selected: `{model_name}`') - -# Initialize Plugins -plugins = [] -if api_keys.get('GEMINI_API_KEY'): - plugins.append(GoogleAI()) -if api_keys.get('ANTHROPIC_API_KEY'): - plugins.append(Anthropic()) -if api_keys.get('DEEPSEEK_API_KEY'): - plugins.append(DeepSeek()) -if api_keys.get('XAI_API_KEY'): - plugins.append(XAI()) -if api_keys.get('OPENAI_API_KEY'): - plugins.append(OpenAI()) - -# Always add Ollama if it has models or just by default (no auth needed typically) -if provider == 'Ollama': - plugins.append(Ollama()) - -if not plugins and provider != 'Ollama': - st.error('No API keys provided for selected provider.') - st.stop() - -# Initialize Genkit (fresh per run) -ai = Genkit( - plugins=plugins, - model=model_name, -) - -# Recover Session Object -session = None -if st.session_state['session_data']: - # Use InMemoryStore with saved data to mimic persistence - # In a real app, this would be a persistent store (Firestore, Redis) - data = st.session_state['session_data'] - - # Reconstruct messages with correct role mapping - msgs = [] - for msg_dict in data['messages']: - role = 'model' if msg_dict['role'] == 'assistant' else msg_dict['role'] - msgs.append(Message(role=role, content=[Part(root=TextPart(text=msg_dict['content']))])) - - # Update data (simplified InMemory structure) - session_store_data: dict[str, Any] = { - data['id']: { - 'id': data['id'], - 'state': data['state'], - 'messages': msgs, - 'created_at': None, - 'updated_at': None, - } - } - - store = InMemorySessionStore(data=session_store_data) - - async def load() -> Session | None: - """Load session.""" - return await ai.load_session(data['id'], store=store) - - session = asyncio.run(load()) - -else: - # No session active - pass - - -def save_session(sess: Session) -> None: - """Save session state to st.session_state.""" - msgs = [] - for m in sess.messages: - # Extract text content (simplification) - text = m.content[0].root.text if m.content and m.content[0].root.text else '' - msgs.append({'role': m.role, 'content': text}) - - st.session_state['session_data'] = { - 'id': sess.id, - 'state': sess.state, - 'messages': msgs, - } - - -def add_log(role: str, text: str) -> None: - """Add a log entry to the session state.""" - st.session_state['logs'].append({'role': role, 'text': text}) - - -async def run_chat(sess: Session, prompt: str) -> GenerateResponseWrapper: - """Run chat asynchronously using session.chat() -> Chat -> send().""" - chat = sess.chat() # Returns a Chat object (matching JS API) - return await chat.send(prompt) - - -# --- Layout --- - -col1, col2 = st.columns([2, 1]) - -with col1: - st.subheader('Interactive Session Flow') - - # Render logs - for log in st.session_state['logs']: - with st.chat_message(log['role']): - st.markdown(log['text']) - - # Step Functions - - # Step 0: Start - if st.session_state['step'] == 0: - if st.button('Start Step 1: Create Session', type='primary'): - sess = ai.create_session(initial_state={'user_name': 'Alice', 'preference': 'concise'}) - save_session(sess) - add_log('system', "Session Created with initial state: `{'user_name': 'Alice', 'preference': 'concise'}`") - st.session_state['step'] = 1 - st.rerun() - - # Step 1: Chat 1 - elif st.session_state['step'] == 1: - if st.button('Run Step 1 Chat', type='primary'): - prompt = "Hi, I'm Alice. What's my name and how many letters are in it?" - add_log('user', prompt) - - with st.spinner('Generating...'): - if session is None: - st.error('Session missing.') - else: - resp = asyncio.run(run_chat(session, prompt)) - save_session(session) - add_log('model', resp.text) - - st.session_state['step'] = 2 - st.rerun() - - # Step 2: Update State - elif st.session_state['step'] == 2: - st.info('Next: Update session state to change topic.') - if st.button('Run Step 2: Update State (Math)', type='primary'): - if session is None: - st.error('Session missing.') - else: - session.update_state({'topic': 'math'}) - save_session(session) - add_log('system', "State Updated: `{'topic': 'math'}`") - st.session_state['step'] = 3 - st.rerun() - - # Step 3: Chat Math - elif st.session_state['step'] == 3: - if st.button('Run Step 3 Chat', type='primary'): - prompt = 'Can you give me a simple problem related to my current topic?' - add_log('user', prompt) - - with st.spinner('Generating...'): - if session is None: - st.error('Session missing.') - else: - resp = asyncio.run(run_chat(session, prompt)) - save_session(session) - add_log('model', resp.text) - - st.session_state['step'] = 4 - st.rerun() - - # Step 4: Update State History - elif st.session_state['step'] == 4: - st.info('Next: Update session state to change topic again.') - if st.button('Run Step 4: Update State (History)', type='primary'): - if session is None: - st.error('Session missing.') - else: - session.update_state({'topic': 'history'}) - save_session(session) - add_log('system', "State Updated: `{'topic': 'history'}`") - st.session_state['step'] = 5 - st.rerun() - - # Step 5: Chat History - elif st.session_state['step'] == 5: - if st.button('Run Step 5 Chat', type='primary'): - prompt = 'Now tell me a fun fact about this new topic.' - add_log('user', prompt) - - with st.spinner('Generating...'): - if session is None: - st.error('Session missing.') - else: - resp = asyncio.run(run_chat(session, prompt)) - save_session(session) - add_log('model', resp.text) - - st.session_state['step'] = 6 - st.rerun() - - elif st.session_state['step'] == 6: - st.success('Demo Complete!') - if st.button('Restart Demo'): - for key in ['step', 'session_data', 'logs']: - del st.session_state[key] - st.rerun() - -with col2: - st.subheader('Current Session State') - if st.session_state['session_data']: - st.json(st.session_state['session_data']['state']) - else: - st.info('No session active.') From 1c7f36dc62a6d606c900d88563b5cb9d123cc170 Mon Sep 17 00:00:00 2001 From: Jeff Huang Date: Thu, 29 Jan 2026 17:30:07 -0600 Subject: [PATCH 2/4] Fix indent --- py/packages/genkit/src/genkit/ai/_registry.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/py/packages/genkit/src/genkit/ai/_registry.py b/py/packages/genkit/src/genkit/ai/_registry.py index a8200d6cb1..4fd19c5254 100644 --- a/py/packages/genkit/src/genkit/ai/_registry.py +++ b/py/packages/genkit/src/genkit/ai/_registry.py @@ -51,9 +51,8 @@ from genkit.aio import ensure_async from genkit.blocks.interfaces import Input, Output - from genkit.blocks.prompt import ExecutablePrompt - from genkit.blocks.resource import FlexibleResourceFn, ResourceOptions - +from genkit.blocks.prompt import ExecutablePrompt +from genkit.blocks.resource import FlexibleResourceFn, ResourceOptions from genkit.blocks.embedding import EmbedderFn, EmbedderOptions from genkit.blocks.evaluator import BatchEvaluatorFn, EvaluatorFn from genkit.blocks.formats.types import FormatDef From 463f6c50d143f30a6a60fed2993b768c3dcf2be6 Mon Sep 17 00:00:00 2001 From: Jeff Huang Date: Thu, 29 Jan 2026 18:06:03 -0600 Subject: [PATCH 3/4] Fix ruff --- py/packages/genkit/src/genkit/ai/_aio.py | 4 +- py/packages/genkit/src/genkit/ai/_registry.py | 7 +- .../genkit/src/genkit/blocks/interfaces.py | 17 ++-- .../genkit/src/genkit/blocks/prompt.py | 3 +- .../genkit/src/genkit/blocks/session/types.py | 18 ++-- py/uv.lock | 88 ------------------- 6 files changed, 29 insertions(+), 108 deletions(-) diff --git a/py/packages/genkit/src/genkit/ai/_aio.py b/py/packages/genkit/src/genkit/ai/_aio.py index 62f9fb1492..8e8cffbd61 100644 --- a/py/packages/genkit/src/genkit/ai/_aio.py +++ b/py/packages/genkit/src/genkit/ai/_aio.py @@ -47,14 +47,13 @@ StreamingCallback as ModelStreamingCallback, generate_action, ) +from genkit.blocks.interfaces import Input as _Input, Output, OutputConfigDict from genkit.blocks.model import ( GenerateResponseChunkWrapper, GenerateResponseWrapper, ModelMiddleware, ) from genkit.blocks.prompt import PromptConfig, load_prompt_folder, to_generate_action_options -from genkit.blocks.interfaces import Input as _Input -from genkit.blocks.interfaces import Output, OutputConfigDict from genkit.blocks.retriever import IndexerRef, IndexerRequest, RetrieverRef from genkit.core.action import Action, ActionRunContext from genkit.core.action.types import ActionKind @@ -141,7 +140,6 @@ def _resolve_embedder_name(self, embedder: str | EmbedderRef | None) -> str: else: raise ValueError('Embedder must be specified as a string name or an EmbedderRef.') - @overload async def generate( self, diff --git a/py/packages/genkit/src/genkit/ai/_registry.py b/py/packages/genkit/src/genkit/ai/_registry.py index 4fd19c5254..8158012785 100644 --- a/py/packages/genkit/src/genkit/ai/_registry.py +++ b/py/packages/genkit/src/genkit/ai/_registry.py @@ -49,15 +49,13 @@ from typing_extensions import Never, TypeVar from genkit.aio import ensure_async - -from genkit.blocks.interfaces import Input, Output -from genkit.blocks.prompt import ExecutablePrompt -from genkit.blocks.resource import FlexibleResourceFn, ResourceOptions from genkit.blocks.embedding import EmbedderFn, EmbedderOptions from genkit.blocks.evaluator import BatchEvaluatorFn, EvaluatorFn from genkit.blocks.formats.types import FormatDef +from genkit.blocks.interfaces import Input, Output from genkit.blocks.model import ModelFn, ModelMiddleware from genkit.blocks.prompt import ( + ExecutablePrompt, define_helper, define_partial, define_prompt, @@ -71,6 +69,7 @@ define_reranker as define_reranker_block, rerank as rerank_block, ) +from genkit.blocks.resource import FlexibleResourceFn, ResourceOptions from genkit.blocks.retriever import IndexerFn, RetrieverFn from genkit.blocks.tools import ToolRunContext from genkit.codec import dump_dict diff --git a/py/packages/genkit/src/genkit/blocks/interfaces.py b/py/packages/genkit/src/genkit/blocks/interfaces.py index 54c28cc587..d24d8fff7a 100644 --- a/py/packages/genkit/src/genkit/blocks/interfaces.py +++ b/py/packages/genkit/src/genkit/blocks/interfaces.py @@ -18,7 +18,7 @@ from __future__ import annotations -from typing import Any, Generic, Protocol, TypedDict, TypeVar +from typing import Generic, Protocol, TypedDict, TypeVar InputT = TypeVar('InputT') OutputT = TypeVar('OutputT') @@ -48,7 +48,6 @@ def __init__(self, schema: type[InputT]) -> None: Args: schema: The type/class for structured input. """ - self.schema: type[InputT] = schema @@ -77,7 +76,6 @@ def __init__( instructions: Optional formatting instructions. constrained: Whether to constrain output to schema. """ - self.schema: type[OutputT] = schema self.format: str = format self.content_type: str | None = content_type @@ -86,7 +84,6 @@ def __init__( def to_dict(self) -> OutputConfigDict: """Convert to OutputConfigDict for internal use.""" - result: OutputConfigDict = {'schema': self.schema, 'format': self.format} if self.content_type is not None: result['content_type'] = self.content_type @@ -100,8 +97,14 @@ def to_dict(self) -> OutputConfigDict: class ExecutablePromptLike(Protocol): """Minimal interface for prompt-like objects used in typing.""" - async def __call__(self, *args: Any, **kwargs: Any) -> Any: ... + async def __call__(self, *args: object, **kwargs: object) -> object: + """Execute the prompt with positional and keyword arguments.""" + ... - def stream(self, *args: Any, **kwargs: Any) -> Any: ... + def stream(self, *args: object, **kwargs: object) -> object: + """Stream prompt execution results.""" + ... - async def render(self, *args: Any, **kwargs: Any) -> Any: ... + async def render(self, *args: object, **kwargs: object) -> object: + """Render the prompt into a concrete representation.""" + ... diff --git a/py/packages/genkit/src/genkit/blocks/prompt.py b/py/packages/genkit/src/genkit/blocks/prompt.py index 3eba11f5b3..b4f794b884 100644 --- a/py/packages/genkit/src/genkit/blocks/prompt.py +++ b/py/packages/genkit/src/genkit/blocks/prompt.py @@ -132,7 +132,7 @@ from collections.abc import AsyncIterable, Awaitable, Callable from pathlib import Path from typing import Any, ClassVar, Generic, TypedDict, TypeVar, cast, overload -from genkit.blocks.interfaces import Input, Output + from dotpromptz.typing import ( DataArgument, PromptFunction, @@ -147,6 +147,7 @@ generate_action, to_tool_definition, ) +from genkit.blocks.interfaces import Input, Output from genkit.blocks.model import ( GenerateResponseChunkWrapper, GenerateResponseWrapper, diff --git a/py/packages/genkit/src/genkit/blocks/session/types.py b/py/packages/genkit/src/genkit/blocks/session/types.py index b08ae2e793..74f3c59323 100644 --- a/py/packages/genkit/src/genkit/blocks/session/types.py +++ b/py/packages/genkit/src/genkit/blocks/session/types.py @@ -18,7 +18,7 @@ from __future__ import annotations -from typing import Any, Protocol +from typing import Protocol from genkit.types import Message @@ -26,7 +26,9 @@ class GenkitLike(Protocol): """Minimal Genkit surface needed by session/chat.""" - async def generate(self, *args: Any, **kwargs: Any) -> Any: ... + async def generate(self, *args: object, **kwargs: object) -> object: + """Generate content using the Genkit instance.""" + ... class SessionLike(Protocol): @@ -35,8 +37,14 @@ class SessionLike(Protocol): id: str _ai: GenkitLike - def get_messages(self, thread_name: str) -> list[Message]: ... + def get_messages(self, thread_name: str) -> list[Message]: + """Return messages for the named thread.""" + ... - def update_messages(self, thread_name: str, messages: list[Message]) -> None: ... + def update_messages(self, thread_name: str, messages: list[Message]) -> None: + """Replace messages for the named thread.""" + ... - async def save(self) -> None: ... + async def save(self) -> None: + """Persist session state.""" + ... diff --git a/py/uv.lock b/py/uv.lock index 356ee24ebd..5a14ce0275 100644 --- a/py/uv.lock +++ b/py/uv.lock @@ -12,7 +12,6 @@ resolution-markers = [ [manifest] members = [ "anthropic-hello", - "chat-demo", "compat-oai-hello", "deepseek-hello", "dev-local-vectorstore-hello", @@ -48,7 +47,6 @@ members = [ "ollama-simple-embed", "prompt-demo", "realtime-tracing-demo", - "session-demo", "short-n-long", "tool-interrupts", "vertex-ai-vector-search-bigquery", @@ -772,52 +770,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0a/4c/925909008ed5a988ccbb72dcc897407e5d6d3bd72410d69e051fc0c14647/charset_normalizer-3.4.4-py3-none-any.whl", hash = "sha256:7a32c560861a02ff789ad905a2fe94e3f840803362c84fecf1851cb4cf3dc37f", size = 53402, upload-time = "2025-10-14T04:42:31.76Z" }, ] -[[package]] -name = "chat-demo" -version = "0.1.0" -source = { editable = "samples/chat-demo" } -dependencies = [ - { name = "genkit" }, - { name = "genkit-plugin-anthropic" }, - { name = "genkit-plugin-compat-oai" }, - { name = "genkit-plugin-deepseek" }, - { name = "genkit-plugin-google-cloud" }, - { name = "genkit-plugin-google-genai" }, - { name = "genkit-plugin-ollama" }, - { name = "genkit-plugin-vertex-ai" }, - { name = "genkit-plugin-xai" }, - { name = "ollama" }, - { name = "pydantic" }, - { name = "streamlit" }, - { name = "structlog" }, - { name = "uvloop" }, -] - -[package.optional-dependencies] -dev = [ - { name = "watchdog" }, -] - -[package.metadata] -requires-dist = [ - { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-anthropic", editable = "plugins/anthropic" }, - { name = "genkit-plugin-compat-oai", editable = "plugins/compat-oai" }, - { name = "genkit-plugin-deepseek", editable = "plugins/deepseek" }, - { name = "genkit-plugin-google-cloud", editable = "plugins/google-cloud" }, - { name = "genkit-plugin-google-genai", editable = "plugins/google-genai" }, - { name = "genkit-plugin-ollama", editable = "plugins/ollama" }, - { name = "genkit-plugin-vertex-ai", editable = "plugins/vertex-ai" }, - { name = "genkit-plugin-xai", editable = "plugins/xai" }, - { name = "ollama" }, - { name = "pydantic", specifier = ">=2.10.5" }, - { name = "streamlit", specifier = ">=1.40.0" }, - { name = "structlog", specifier = ">=25.2.0" }, - { name = "uvloop", specifier = ">=0.21.0" }, - { name = "watchdog", marker = "extra == 'dev'", specifier = ">=6.0.0" }, -] -provides-extras = ["dev"] - [[package]] name = "click" version = "8.3.1" @@ -6427,46 +6379,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1c/78/504fdd027da3b84ff1aecd9f6957e65f35134534ccc6da8628eb71e76d3f/send2trash-2.1.0-py3-none-any.whl", hash = "sha256:0da2f112e6d6bb22de6aa6daa7e144831a4febf2a87261451c4ad849fe9a873c", size = 17610, upload-time = "2026-01-14T06:27:35.218Z" }, ] -[[package]] -name = "session-demo" -version = "0.1.0" -source = { editable = "samples/session-demo" } -dependencies = [ - { name = "genkit" }, - { name = "genkit-plugin-anthropic" }, - { name = "genkit-plugin-deepseek" }, - { name = "genkit-plugin-google-genai" }, - { name = "genkit-plugin-ollama" }, - { name = "genkit-plugin-xai" }, - { name = "ollama" }, - { name = "pydantic" }, - { name = "streamlit" }, - { name = "structlog" }, - { name = "uvloop" }, -] - -[package.optional-dependencies] -dev = [ - { name = "watchdog" }, -] - -[package.metadata] -requires-dist = [ - { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-anthropic", editable = "plugins/anthropic" }, - { name = "genkit-plugin-deepseek", editable = "plugins/deepseek" }, - { name = "genkit-plugin-google-genai", editable = "plugins/google-genai" }, - { name = "genkit-plugin-ollama", editable = "plugins/ollama" }, - { name = "genkit-plugin-xai", editable = "plugins/xai" }, - { name = "ollama" }, - { name = "pydantic", specifier = ">=2.10.5" }, - { name = "streamlit", specifier = ">=1.40.0" }, - { name = "structlog", specifier = ">=25.2.0" }, - { name = "uvloop", specifier = ">=0.21.0" }, - { name = "watchdog", marker = "extra == 'dev'", specifier = ">=6.0.0" }, -] -provides-extras = ["dev"] - [[package]] name = "setuptools" version = "80.10.1" From b537c3ddc4307933f6dc3ea350f7e31a01e5f67f Mon Sep 17 00:00:00 2001 From: Jeff Huang Date: Thu, 29 Jan 2026 20:25:14 -0600 Subject: [PATCH 4/4] fix ty checks --- .../genkit/src/genkit/blocks/interfaces.py | 18 +------- .../genkit/src/genkit/blocks/session/chat.py | 46 +++++++++++++++---- .../src/genkit/blocks/session/session.py | 21 +++++---- .../genkit/src/genkit/blocks/session/types.py | 20 +++++++- 4 files changed, 68 insertions(+), 37 deletions(-) diff --git a/py/packages/genkit/src/genkit/blocks/interfaces.py b/py/packages/genkit/src/genkit/blocks/interfaces.py index d24d8fff7a..2c24d49e1e 100644 --- a/py/packages/genkit/src/genkit/blocks/interfaces.py +++ b/py/packages/genkit/src/genkit/blocks/interfaces.py @@ -18,7 +18,7 @@ from __future__ import annotations -from typing import Generic, Protocol, TypedDict, TypeVar +from typing import Generic, TypedDict, TypeVar InputT = TypeVar('InputT') OutputT = TypeVar('OutputT') @@ -92,19 +92,3 @@ def to_dict(self) -> OutputConfigDict: if self.constrained is not None: result['constrained'] = self.constrained return result - - -class ExecutablePromptLike(Protocol): - """Minimal interface for prompt-like objects used in typing.""" - - async def __call__(self, *args: object, **kwargs: object) -> object: - """Execute the prompt with positional and keyword arguments.""" - ... - - def stream(self, *args: object, **kwargs: object) -> object: - """Stream prompt execution results.""" - ... - - async def render(self, *args: object, **kwargs: object) -> object: - """Render the prompt into a concrete representation.""" - ... diff --git a/py/packages/genkit/src/genkit/blocks/session/chat.py b/py/packages/genkit/src/genkit/blocks/session/chat.py index 412edc58e8..c5f43c5ef0 100644 --- a/py/packages/genkit/src/genkit/blocks/session/chat.py +++ b/py/packages/genkit/src/genkit/blocks/session/chat.py @@ -111,7 +111,7 @@ from genkit.aio import Channel from genkit.blocks.model import GenerateResponseChunkWrapper, GenerateResponseWrapper -from genkit.types import Message, Part +from genkit.types import Message, Part, TextPart from .store import MAIN_THREAD from .types import SessionLike @@ -196,6 +196,18 @@ def response(self) -> asyncio.Future[GenerateResponseWrapper]: return self._response_future +def _normalize_prompt_parts(prompt: str | Part | list[Part]) -> list[Part]: + """Normalize prompt input into a list of Parts.""" + if isinstance(prompt, str): + # Part is a RootModel, so we pass content via the 'root' parameter. + return [Part(root=TextPart(text=prompt))] + if isinstance(prompt, list): + return list(prompt) + if isinstance(prompt, Part): # pyright: ignore[reportUnnecessaryIsInstance] + return [prompt] + return [] + + class Chat: """A stateful chat conversation. @@ -298,7 +310,7 @@ async def send( ) -> GenerateResponseWrapper: """Send a message and get the complete response (matches JS chat.send).""" # Append user message to history - self._messages.append(Message(role='user', content=[Part.from_text(prompt)])) + self._messages.append(Message(role='user', content=_normalize_prompt_parts(prompt))) # Merge base options and call-specific options gen_options = {**self._request_base, **kwargs} @@ -308,7 +320,14 @@ async def send( response = await self._session._ai.generate(**gen_options) # pyright: ignore[reportPrivateUsage] # Append model response to history - self._messages.append(Message(role='assistant', content=response.content)) + if response.message is not None: + self._messages.append( + Message( + role=response.message.role, + content=response.message.content, + metadata=response.message.metadata, + ) + ) self._session.update_messages(self._thread_name, self._messages) await self._session.save() @@ -321,25 +340,34 @@ def send_stream( ) -> ChatStreamResponse: """Send a message and stream the response (matches JS chat.sendStream).""" # Append user message to history - self._messages.append(Message(role='user', content=[Part.from_text(prompt)])) + self._messages.append(Message(role='user', content=_normalize_prompt_parts(prompt))) # Merge base options and call-specific options gen_options = {**self._request_base, **kwargs} gen_options['messages'] = self._messages # Use session's AI instance to generate response (streaming) - stream_result = self._session._ai.generate_stream(**gen_options) # pyright: ignore[reportPrivateUsage] + stream, base_response_future = self._session._ai.generate_stream( # pyright: ignore[reportPrivateUsage] + **gen_options + ) async def _run_stream() -> GenerateResponseWrapper: - response = await stream_result.response + response = await base_response_future # Append model response to history - self._messages.append(Message(role='assistant', content=response.content)) + if response.message is not None: + self._messages.append( + Message( + role=response.message.role, + content=response.message.content, + metadata=response.message.metadata, + ) + ) self._session.update_messages(self._thread_name, self._messages) await self._session.save() return response - response_future = asyncio.ensure_future(_run_stream()) - return ChatStreamResponse(stream_result.stream, response_future) + wrapped_response_future = asyncio.ensure_future(_run_stream()) + return ChatStreamResponse(stream, wrapped_response_future) def _merge_chat_options(base: dict[str, Any], updates: dict[str, Any]) -> dict[str, Any]: diff --git a/py/packages/genkit/src/genkit/blocks/session/session.py b/py/packages/genkit/src/genkit/blocks/session/session.py index 61c917c4e7..125d2a1a29 100644 --- a/py/packages/genkit/src/genkit/blocks/session/session.py +++ b/py/packages/genkit/src/genkit/blocks/session/session.py @@ -39,7 +39,7 @@ import uuid from typing import Any -from genkit.blocks.interfaces import ExecutablePromptLike +from genkit.blocks.prompt import ExecutablePrompt from genkit.types import Message from .chat import Chat, ChatOptions @@ -118,10 +118,16 @@ def __init__( if data: self._id: str = data.get('id', str(uuid.uuid4())) self._state: dict[str, Any] = data.get('state', {}) or {} - threads = data.get('threads', {}) or {} - if not threads and data.get('messages'): + threads: dict[str, list[Message]] = {} + threads_value = data.get('threads') + if isinstance(threads_value, dict): + threads = { + str(name): messages for name, messages in threads_value.items() if isinstance(messages, list) + } + messages = data.get('messages') + if not threads and messages: # Backwards compatibility: use legacy messages field as main thread - threads = {MAIN_THREAD: data.get('messages', [])} + threads = {MAIN_THREAD: messages} self._threads: dict[str, list[Message]] = threads self._created_at: float = data.get('created_at') or time.time() self._updated_at: float = data.get('updated_at') or self._created_at @@ -190,8 +196,8 @@ async def save(self) -> None: def chat( self, - thread_or_preamble_or_options: str | ExecutablePromptLike | ChatOptions | None = None, - preamble_or_options: ExecutablePromptLike | ChatOptions | None = None, + thread_or_preamble_or_options: str | ExecutablePrompt | ChatOptions | None = None, + preamble_or_options: ExecutablePrompt | ChatOptions | None = None, options: ChatOptions | None = None, ) -> Chat: """Create a Chat object for multi-turn conversation (matching JS API). @@ -279,7 +285,4 @@ def chat( def _is_executable_prompt(self, obj: Any) -> bool: # noqa: ANN401 """Check if an object is an ExecutablePrompt.""" - # Import here to avoid circular imports - from genkit.blocks.prompt import ExecutablePrompt - return isinstance(obj, ExecutablePrompt) diff --git a/py/packages/genkit/src/genkit/blocks/session/types.py b/py/packages/genkit/src/genkit/blocks/session/types.py index 74f3c59323..6af48b5a66 100644 --- a/py/packages/genkit/src/genkit/blocks/session/types.py +++ b/py/packages/genkit/src/genkit/blocks/session/types.py @@ -18,25 +18,41 @@ from __future__ import annotations +import asyncio from typing import Protocol +from genkit.aio import Channel +from genkit.blocks.model import GenerateResponseChunkWrapper, GenerateResponseWrapper from genkit.types import Message class GenkitLike(Protocol): """Minimal Genkit surface needed by session/chat.""" - async def generate(self, *args: object, **kwargs: object) -> object: + async def generate(self, *args: object, **kwargs: object) -> GenerateResponseWrapper[object]: """Generate content using the Genkit instance.""" ... + def generate_stream( + self, *args: object, **kwargs: object + ) -> tuple[ + Channel[GenerateResponseChunkWrapper, GenerateResponseWrapper[object]], + asyncio.Future[GenerateResponseWrapper[object]], + ]: + """Generate streaming content using the Genkit instance.""" + ... + class SessionLike(Protocol): """Minimal Session surface needed by Chat.""" - id: str _ai: GenkitLike + @property + def id(self) -> str: + """Return the session ID.""" + ... + def get_messages(self, thread_name: str) -> list[Message]: """Return messages for the named thread.""" ...