Skip to content

refactor(py): move session/chat into blocks and remove public API #4321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
huangjeff5 merged 4 commits into from
Jan 30, 2026
Merged

refactor(py): move session/chat into blocks and remove public API #4321

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0169dc5
refactor(py): move session/chat into blocks and remove public API
huangjeff5 Jan 29, 2026
1c7f36d
Fix indent
huangjeff5 Jan 29, 2026
463f6c5
Fix ruff
huangjeff5 Jan 30, 2026
b537c3d
fix ty checks
huangjeff5 Jan 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 0 additions & 6 deletions py/packages/genkit/src/genkit/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,9 +38,6 @@ async def hello(name: str) -> str:
GENKIT_VERSION,
ActionKind,
ActionRunContext,
Chat,
ChatOptions,
ChatStreamResponse,
ExecutablePrompt,
FlowWrapper,
GenerateStreamResponse,
Expand Down Expand Up @@ -101,9 +98,6 @@ async def hello(name: str) -> str:
# From genkit.ai
'ActionKind',
'ActionRunContext',
'Chat',
'ChatOptions',
'ChatStreamResponse',
'ExecutablePrompt',
'FlowWrapper',
'GenerateStreamResponse',
Expand Down
5 changes: 0 additions & 5 deletions py/packages/genkit/src/genkit/ai/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -61,10 +60,6 @@
'OutputOptions',
'PromptGenerateOptions',
'ResumeOptions',
# Session/Chat
'Chat',
'ChatOptions',
'ChatStreamResponse',
# Document
'Document',
# Plugin
Expand Down
298 changes: 3 additions & 295 deletions py/packages/genkit/src/genkit/ai/_aio.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -50,6 +47,7 @@
StreamingCallback as ModelStreamingCallback,
generate_action,
)
from genkit.blocks.interfaces import Input as _Input, Output, OutputConfigDict
from genkit.blocks.model import (
GenerateResponseChunkWrapper,
GenerateResponseWrapper,
Expand All @@ -70,7 +68,6 @@
Operation,
SpanMetadata,
)
from genkit.session import Chat, ChatOptions, InMemorySessionStore, Session, SessionStore
from genkit.types import (
DocumentData,
GenerationCommonConfig,
Expand All @@ -86,135 +83,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."""

Expand Down Expand Up @@ -265,173 +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.')

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(
self,
Expand Down
Loading
Loading