Functional API

import asyncio

from dotenv import load_dotenv

import synalinks

# =============================================================================
# STEP 1: Define Your Data Models
# =============================================================================
# Data Models are like forms that define what information your program
# expects as input and what it will produce as output.
#
# Each field has:
#   - A name (e.g., "query")
#   - A type (e.g., str for text)
#   - A description (helps the LLM understand what to fill in)


class Query(synalinks.DataModel):
    """The input to our program - a user's question."""

    query: str = synalinks.Field(
        description="The user query",
    )


class AnswerWithThinking(synalinks.DataModel):
    """The output from our program - reasoning + final answer.

    By asking the LLM to show its thinking, we get better answers.
    This is called "Chain of Thought" prompting.
    """

    thinking: str = synalinks.Field(
        description="Your step by step thinking",
    )
    answer: str = synalinks.Field(
        description="The correct answer",
    )


# =============================================================================
# STEP 2: Build and Run the Program
# =============================================================================


async def main():
    # Load environment variables (like your API key) from .env file
    load_dotenv()

    # Enable observability for tracing (view traces at http://localhost:5000)
#     synalinks.enable_observability(
#         tracking_uri="http://localhost:5000",
#         experiment_name="lesson_1a_functional_api",
#     )

    # -------------------------------------------------------------------------
    # 2.1: Configure the Language Model
    # -------------------------------------------------------------------------
    # The LanguageModel is the AI brain that will process our requests.
    # We're using Google's Gemini model here.
    language_model = synalinks.LanguageModel(
        model="ollama/qwen3:8b",
    )

    # -------------------------------------------------------------------------
    # 2.2: Build the Program with the Functional API
    # -------------------------------------------------------------------------
    print("=" * 60)
    print("Building a Chain-of-Thought Program")
    print("=" * 60)

    # Create the input placeholder
    # This tells Synalinks what kind of data will enter the program
    inputs = synalinks.Input(data_model=Query)

    # Create a Generator module and connect it to the input
    # The Generator uses the LLM to transform input into output
    outputs = await synalinks.Generator(
        data_model=AnswerWithThinking,  # What to produce
        language_model=language_model,  # Which AI to use
    )(inputs)  # <-- This connects the module to our input

    # Wrap everything into a Program
    # Now we have a reusable AI application!
    program = synalinks.Program(
        inputs=inputs,
        outputs=outputs,
        name="chain_of_thought",
        description="Useful to answer in a step by step manner.",
    )

    # Generate a visualization of our program (optional but helpful!)
    synalinks.utils.plot_program(
        program,
        to_folder="examples",
        show_module_names=True,
        show_trainable=True,
        show_schemas=True,
    )

    # -------------------------------------------------------------------------
    # 2.3: Run the Program
    # -------------------------------------------------------------------------
    print("\nRunning the program...")
    print("-" * 60)

    # Call the program like a function, passing in our query
    result = await program(
        Query(query="What are the key aspects of human cognition?"),
    )

    # Display the result in a nicely formatted way
    print("\nResult:")
    print(result.prettify_json())

    # -------------------------------------------------------------------------
    # 2.4: Understanding the Output
    # -------------------------------------------------------------------------
    # The result is a JsonDataModel containing:
    #   - thinking: The LLM's step-by-step reasoning
    #   - answer: The final answer
    #
    # JsonDataModel supports direct field access - no need for get_json()!
    print("\n" + "=" * 60)
    print("Accessing individual fields:")
    print("=" * 60)
    print(f"\nThinking: {result['thinking'][:100]}...")
    print(f"\nAnswer: {result['answer'][:100]}...")


if __name__ == "__main__":
    asyncio.run(main())