API Reference

def __init__(
    self,
    plugins: list[Plugin] | None = None,
    model: str | None = None,
    reflection_server_spec: ServerSpec | None = None,
) -> None:
    """Initialize a new Genkit instance.

    Args:
        plugins: List of plugins to initialize.
        model: Model name to use.
        reflection_server_spec: Server spec for the reflection
            server.
    """
    super().__init__(plugins=plugins, model=model, reflection_server_spec=reflection_server_spec)

def define_batch_evaluator(
    self,
    name: str,
    display_name: str,
    definition: str,
    fn: BatchEvaluatorFn,
    is_billed: bool = False,
    config_schema: BaseModel | dict[str, Any] | None = None,
    metadata: dict[str, Any] | None = None,
    description: str | None = None,
) -> Callable[[Callable], Callable]:
    """Define a batch evaluator action.

    This action runs the callback function on the entire dataset.

    Args:
        name: Name of the evaluator.
        fn: Function implementing the evaluator behavior.
        display_name: User-visible display name
        definition: User-visible evaluator definition
        is_billed: Whether the evaluator performs any billed actions
                    (paid  APIs, LLMs etc.)
        config_schema: Optional schema for evaluator configuration.
        metadata: Optional metadata for the evaluator.
        description: Optional description for the evaluator.
    """
    evaluator_meta = metadata if metadata else {}
    if 'evaluator' not in evaluator_meta:
        evaluator_meta['evaluator'] = {}
    evaluator_meta['evaluator'][EVALUATOR_METADATA_KEY_DEFINITION] = definition
    evaluator_meta['evaluator'][EVALUATOR_METADATA_KEY_DISPLAY_NAME] = display_name
    evaluator_meta['evaluator'][EVALUATOR_METADATA_KEY_IS_BILLED] = is_billed
    if 'label' not in evaluator_meta['evaluator'] or not evaluator_meta['evaluator']['label']:
        evaluator_meta['evaluator']['label'] = name
    if config_schema:
        evaluator_meta['evaluator']['customOptions'] = to_json_schema(config_schema)

    evaluator_description = get_func_description(fn, description)
    return self.registry.register_action(
        name=name,
        kind=ActionKind.EVALUATOR,
        fn=fn,
        metadata=evaluator_meta,
        description=evaluator_description,
    )

def define_embedder(
    self,
    name: str,
    fn: EmbedderFn,
    metadata: dict[str, Any] | None = None,
    description: str | None = None,
) -> Action:
    """Define a custom embedder action.

    Args:
        name: Name of the model.
        fn: Function implementing the embedder behavior.
        metadata: Optional metadata for the model.
        description: Optional description for the embedder.
    """
    embedder_description = get_func_description(fn, description)
    return self.registry.register_action(
        name=name,
        kind=ActionKind.EMBEDDER,
        fn=fn,
        metadata=metadata,
        description=embedder_description,
    )

def define_evaluator(
    self,
    name: str,
    display_name: str,
    definition: str,
    fn: EvaluatorFn,
    is_billed: bool = False,
    config_schema: BaseModel | dict[str, Any] | None = None,
    metadata: dict[str, Any] | None = None,
    description: str | None = None,
) -> Callable[[Callable], Callable]:
    """Define a evaluator action.

    This action runs the callback function on the every sample of
    the input dataset.

    Args:
        name: Name of the evaluator.
        fn: Function implementing the evaluator behavior.
        display_name: User-visible display name
        definition: User-visible evaluator definition
        is_billed: Whether the evaluator performs any billed actions
                    (paid  APIs, LLMs etc.)
        config_schema: Optional schema for evaluator configuration.
        metadata: Optional metadata for the evaluator.
        description: Optional description for the evaluator.
    """
    evaluator_meta = metadata if metadata else {}
    if 'evaluator' not in evaluator_meta:
        evaluator_meta['evaluator'] = {}
    evaluator_meta['evaluator'][EVALUATOR_METADATA_KEY_DEFINITION] = definition
    evaluator_meta['evaluator'][EVALUATOR_METADATA_KEY_DISPLAY_NAME] = display_name
    evaluator_meta['evaluator'][EVALUATOR_METADATA_KEY_IS_BILLED] = is_billed
    if 'label' not in evaluator_meta['evaluator'] or not evaluator_meta['evaluator']['label']:
        evaluator_meta['evaluator']['label'] = name
    if config_schema:
        evaluator_meta['evaluator']['customOptions'] = to_json_schema(config_schema)

    evaluator_description = get_func_description(fn, description)

    def eval_stepper_fn(req: EvalRequest) -> EvalResponse:
        eval_responses: list[EvalFnResponse] = []
        for index in range(len(req.dataset)):
            datapoint = req.dataset[index]
            if datapoint.test_case_id is None:
                datapoint.test_case_id = str(uuid.uuid4())
            span_metadata = SpanMetadata(
                name=f'Test Case {datapoint.test_case_id}',
                metadata={'evaluator:evalRunId': req.eval_run_id},
            )
            try:
                with run_in_new_span(span_metadata, labels={'genkit:type': 'evaluator'}) as span:
                    span_id = span.span_id()
                    trace_id = span.trace_id()
                    try:
                        span.set_input(datapoint)
                        test_case_output = fn(datapoint, req.options)
                        test_case_output.span_id = span_id
                        test_case_output.trace_id = trace_id
                        span.set_output(test_case_output)
                        eval_responses.append(test_case_output)
                    except Exception as e:
                        logger.debug(f'eval_stepper_fn error: {str(e)}')
                        logger.debug(traceback.format_exc())
                        evaluation = Score(
                            error=f'Evaluation of test case {datapoint.test_case_id} failed: \n{str(e)}'
                        )
                        eval_responses.append(
                            EvalFnResponse(
                                span_id=span_id,
                                trace_id=trace_id,
                                test_case_id=datapoint.test_case_id,
                                evaluation=evaluation,
                            )
                        )
                        # Raise to mark span as failed
                        raise e
            except Exception:
                # Continue to process other points
                continue
        return EvalResponse(eval_responses)

    return self.registry.register_action(
        name=name,
        kind=ActionKind.EVALUATOR,
        fn=eval_stepper_fn,
        metadata=evaluator_meta,
        description=evaluator_description,
    )

def define_format(self, format: FormatDef) -> None:
    """Registers a custom format in the registry.

    Args:
        format: The format to register.
    """
    self.registry.register_value('format', format.name, format)

def define_model(
    self,
    name: str,
    fn: ModelFn,
    config_schema: BaseModel | dict[str, Any] | None = None,
    metadata: dict[str, Any] | None = None,
    info: ModelInfo | None = None,
    description: str | None = None,
) -> Action:
    """Define a custom model action.

    Args:
        name: Name of the model.
        fn: Function implementing the model behavior.
        config_schema: Optional schema for model configuration.
        metadata: Optional metadata for the model.
        info: Optional ModelInfo for the model.
        description: Optional description for the model.
    """
    model_meta: dict[str, Any] = metadata if metadata else {}
    if info:
        model_meta['model'] = dump_dict(info)
    if 'model' not in model_meta:
        model_meta['model'] = {}
    if 'label' not in model_meta['model'] or not model_meta['model']['label']:
        model_meta['model']['label'] = name

    if config_schema:
        model_meta['model']['customOptions'] = to_json_schema(config_schema)

    model_description = get_func_description(fn, description)
    return self.registry.register_action(
        name=name,
        kind=ActionKind.MODEL,
        fn=fn,
        metadata=model_meta,
        description=model_description,
    )

def define_prompt(
    self,
    variant: str | None = None,
    model: str | None = None,
    config: GenerationCommonConfig | dict[str, Any] | None = None,
    description: str | None = None,
    input_schema: type | dict[str, Any] | None = None,
    system: str | Part | list[Part] | None = None,
    prompt: str | Part | list[Part] | None = None,
    messages: str | list[Message] | None = None,
    output_format: str | None = None,
    output_content_type: str | None = None,
    output_instructions: bool | str | None = None,
    output_schema: type | dict[str, Any] | None = None,
    output_constrained: bool | None = None,
    max_turns: int | None = None,
    return_tool_requests: bool | None = None,
    metadata: dict[str, Any] | None = None,
    tools: list[str] | None = None,
    tool_choice: ToolChoice | None = None,
    use: list[ModelMiddleware] | None = None,
    # TODO:
    #  docs: list[Document]
):
    """Define a prompt.

    Args:
        variant: Optional variant name for the prompt.
        model: Optional model name to use for the prompt.
        config: Optional configuration for the model.
        description: Optional description for the prompt.
        input_schema: Optional schema for the input to the prompt.
        system: Optional system message for the prompt.
        prompt: Optional prompt for the model.
        messages: Optional messages for the model.
        output_format: Optional output format for the prompt.
        output_content_type: Optional output content type for the prompt.
        output_instructions: Optional output instructions for the prompt.
        output_schema: Optional schema for the output from the prompt.
        output_constrained: Optional flag indicating whether the output
            should be constrained.
        max_turns: Optional maximum number of turns for the prompt.
        return_tool_requests: Optional flag indicating whether tool requests
            should be returned.
        metadata: Optional metadata for the prompt.
        tools: Optional list of tools to use for the prompt.
        tool_choice: Optional tool choice for the prompt.
        use: Optional list of model middlewares to use for the prompt.
    """
    return define_prompt(
        self.registry,
        variant=variant,
        model=model,
        config=config,
        description=description,
        input_schema=input_schema,
        system=system,
        prompt=prompt,
        messages=messages,
        output_format=output_format,
        output_content_type=output_content_type,
        output_instructions=output_instructions,
        output_schema=output_schema,
        output_constrained=output_constrained,
        max_turns=max_turns,
        return_tool_requests=return_tool_requests,
        metadata=metadata,
        tools=tools,
        tool_choice=tool_choice,
        use=use,
    )

def define_retriever(
    self,
    name: str,
    fn: RetrieverFn,
    config_schema: BaseModel | dict[str, Any] | None = None,
    metadata: dict[str, Any] | None = None,
    description: str | None = None,
) -> Callable[[Callable], Callable]:
    """Define a retriever action.

    Args:
        name: Name of the retriever.
        fn: Function implementing the retriever behavior.
        config_schema: Optional schema for retriever configuration.
        metadata: Optional metadata for the retriever.
        description: Optional description for the retriever.
    """
    retriever_meta = metadata if metadata else {}
    if 'retriever' not in retriever_meta:
        retriever_meta['retriever'] = {}
    if 'label' not in retriever_meta['retriever'] or not retriever_meta['retriever']['label']:
        retriever_meta['retriever']['label'] = name
    if config_schema:
        retriever_meta['retriever']['customOptions'] = to_json_schema(config_schema)

    retriever_description = get_func_description(fn, description)
    return self.registry.register_action(
        name=name,
        kind=ActionKind.RETRIEVER,
        fn=fn,
        metadata=retriever_meta,
        description=retriever_description,
    )

async def embed(
    self,
    embedder: str | None = None,
    documents: list[Document] | None = None,
    options: dict[str, Any] | None = None,
) -> EmbedResponse:
    """Calculates embeddings for documents.

    Args:
        embedder: Optional embedder model name to use.
        documents: Texts to embed.
        options: embedding options

    Returns:
        The generated response with embeddings.
    """
    embed_action = self.registry.lookup_action(ActionKind.EMBEDDER, embedder)

    return (await embed_action.arun(EmbedRequest(input=documents, options=options))).response

def flow(self, name: str | None = None, description: str | None = None) -> Callable[[Callable], Callable]:
    """Decorator to register a function as a flow.

    Args:
        name: Optional name for the flow. If not provided, uses the
            function name.
        description: Optional description for the flow. If not provided,
            uses the function docstring.

    Returns:
        A decorator function that registers the flow.
    """

    def wrapper(func: Callable) -> Callable:
        """Register the decorated function as a flow.

        Args:
            func: The function to register as a flow.

        Returns:
            The wrapped function that executes the flow.
        """
        flow_name = name if name is not None else func.__name__
        flow_description = get_func_description(func, description)
        action = self.registry.register_action(
            name=flow_name,
            kind=ActionKind.FLOW,
            fn=func,
            description=flow_description,
            span_metadata={'genkit:metadata:flow:name': flow_name},
        )

        @wraps(func)
        async def async_wrapper(*args, **kwargs):
            """Asynchronous wrapper for the flow function.

            Args:
                *args: Positional arguments to pass to the flow function.
                **kwargs: Keyword arguments to pass to the flow function.

            Returns:
                The response from the flow function.
            """
            return (await action.arun(*args, **kwargs)).response

        @wraps(func)
        def sync_wrapper(*args, **kwargs):
            """Synchronous wrapper for the flow function.

            Args:
                *args: Positional arguments to pass to the flow function.
                **kwargs: Keyword arguments to pass to the flow function.

            Returns:
                The response from the flow function.
            """
            return action.run(*args, **kwargs).response

        return FlowWrapper(
            fn=async_wrapper if action.is_async else sync_wrapper,
            action=action,
        )

    return wrapper

async def generate(
    self,
    model: str | None = None,
    prompt: str | Part | list[Part] | None = None,
    system: str | Part | list[Part] | None = None,
    messages: list[Message] | None = None,
    tools: list[str] | None = None,
    return_tool_requests: bool | None = None,
    tool_choice: ToolChoice = None,
    tool_responses: list[Part] | None = None,
    config: GenerationCommonConfig | dict[str, Any] | None = None,
    max_turns: int | None = None,
    on_chunk: ModelStreamingCallback | None = None,
    context: dict[str, Any] | None = None,
    output_format: str | None = None,
    output_content_type: str | None = None,
    output_instructions: bool | str | None = None,
    output_schema: type | dict[str, Any] | None = None,
    output_constrained: bool | None = None,
    use: list[ModelMiddleware] | None = None,
    docs: list[DocumentData] | None = None,
) -> GenerateResponseWrapper:
    """Generates text or structured data using a language model.

    This function provides a flexible interface for interacting with various
    language models, supporting both simple text generation and more complex
    interactions involving tools and structured conversations.

    Args:
        model: Optional. The name of the model to use for generation. If not
            provided, a default model may be used.
        prompt: Optional. A single prompt string, a `Part` object, or a list
            of `Part` objects to provide as input to the model. This is used
            for simple text generation.
        system: Optional. A system message string, a `Part` object, or a
            list of `Part` objects to provide context or instructions to
            the model, especially for chat-based models.
        messages: Optional. A list of `Message` objects representing a
            conversation history.  This is used for chat-based models to
            maintain context.
        tools: Optional. A list of tool names (strings) that the model can
            use.
        return_tool_requests: Optional. If `True`, the model will return
            tool requests instead of executing them directly.
        tool_choice: Optional. A `ToolChoice` object specifying how the
            model should choose which tool to use.
        tool_responses: Optional. tool_responses should contain a list of
            tool response parts corresponding to interrupt tool request
            parts from the most recent model message. Each entry must have
            a matching `name` and `ref` (if supplied) for its tool request
            counterpart.
        config: Optional. A `GenerationCommonConfig` object or a dictionary
            containing configuration parameters for the generation process.
            This allows fine-tuning the model's behavior.
        max_turns: Optional. The maximum number of turns in a conversation.
        on_chunk: Optional. A callback function of type
            `ModelStreamingCallback` that is called for each chunk of
            generated text during streaming.
        context: Optional. A dictionary containing additional context
            information that can be used during generation.
        output_format: Optional. The format to use for the output (e.g.,
            'json').
        output_content_type: Optional. The content type of the output.
        output_instructions: Optional. Instructions for formatting the
            output.
        output_schema: Optional. Schema defining the structure of the
            output.
        output_constrained: Optional. Whether to constrain the output to the
            schema.
        use: Optional. A list of `ModelMiddleware` functions to apply to the
            generation process. Middleware can be used to intercept and
            modify requests and responses.
        docs: Optional. A list of documents to be used for grounding.


    Returns:
        A `GenerateResponseWrapper` object containing the model's response,
        which may include generated text, tool requests, or other relevant
        information.

    Note:
        - The `tools`, `return_tool_requests`, and `tool_choice` arguments
          are used for models that support tool usage.
        - The `on_chunk` argument enables streaming responses, allowing you
          to process the generated content as it becomes available.
    """
    return await generate_action(
        self.registry,
        to_generate_action_options(
            registry=self.registry,
            model=model,
            prompt=prompt,
            system=system,
            messages=messages,
            tools=tools,
            return_tool_requests=return_tool_requests,
            tool_choice=tool_choice,
            tool_responses=tool_responses,
            config=config,
            max_turns=max_turns,
            output_format=output_format,
            output_content_type=output_content_type,
            output_instructions=output_instructions,
            output_schema=output_schema,
            output_constrained=output_constrained,
            docs=docs,
        ),
        on_chunk=on_chunk,
        middleware=use,
        context=context if context else ActionRunContext._current_context(),
    )