A function-guess-game with gemini.

function_guess_game.py

import dspy
import os
import numpy as np
import math
from dotenv import load_dotenv
import matplotlib.pyplot as plt


# --- Environment Setup ---
load_dotenv(".env.google")  # Make sure you have your API key in this file

# Configure the language model (replace with your preferred model if needed)
lm = dspy.LM(
    model="gemini/gemini-2.0-flash",
    api_key=os.environ.get("GEMINI_API_KEY"),
)
dspy.configure(lm=lm)


# --- Define the Hidden Function ---
# We'll select one function for the game instance.
# Add more complex functions as desired!
def get_hidden_function(function_name):
    """Returns the specified function for the game."""
    functions = {
        "Linear": lambda x: 2 * x + 3,
        "Quadratic": lambda x: x**2 - 4,
        "Cubic": lambda x: 0.5 * x**3 - 2 * x - 10,
        "Exponential": lambda x: 2**x if x < 10 else float("inf"),
        "Logarithmic": lambda x: math.log(x + 5) if x > -5 else float("inf"),
        "x*sin(pi*x)": lambda x: x * np.sin(np.pi * x),
        "Sine": lambda x: math.sin(x),
    }
    if function_name in functions:
        print(
            f"DEBUG: Hidden function chosen: {function_name}"
        )  # Keep this for debugging
        return function_name, functions[function_name]
    else:
        raise ValueError(
            f"Function '{function_name}' not found. Available functions: {list(functions.keys())}"
        )


# --- DSPy Signatures for the Function Discovery Game ---


class FunctionProber(dspy.Signature):
    """
    You are an analytical agent trying to discover the properties of an unknown mathematical function, f(x).
    You can query the function by providing an input value 'x' and observing the output 'f(x)'.

    Your goal is to deduce the nature of the function (e.g., linear, polynomial degree, exponential, periodic, intercepts, asymptotes, etc.). You need to do so in a limited number of queries, so choose wisely.

    # Information Provided:
    - Knowledge Summary ("what_you_know_thus_far"): This summarizes your current understanding, including:
        * "Confirmed property": Properties you have deduced with high confidence based on the data.
        * "Ruled-out property": Properties you have refuted with high confidence.
        * "Hypothesis": Potential properties or function types you are currently investigating.
    - Query History ("query_history"): A list of dictionaries showing past queries and their results (e.g., [{'input': 1, 'output': 5}, {'input': 2, 'output': 7}, ...]).

    # Your Task: Decide the best numerical input 'x' for the *next* query so you can learn more about the function.

    # Decision Process:
    1. Analyze the Query History: Look for patterns, trends, and special points (like intercepts, minimums, maximums if apparent) in the (input, output) pairs.
    2. Evaluate Knowledge Summary:
        * Prioritize "Confirmed properties": Does your current understanding suggest a specific area to probe further or confirm consistency?
        * Consider "Ruled-out properties": Ensure your next query isn't based on an idea already disproven.
        * Test Hypotheses: If you have active "Hypotheses", choose an 'x' value that will best *differentiate* between competing hypotheses or provide strong evidence for/against a specific one. For example:
            - To test linearity: Query a third point and check for collinearity with two existing points.
            - To test symmetry (e.g., for even/odd polynomials): Query f(-x) if you have f(x).
            - To check for asymptotes: Query large positive or negative values, or values near suspected discontinuities.
            - To test periodicity: Query x + P if you suspect a period P based on existing points.
            - To determine polynomial degree: Query multiple points to see how differences between outputs change.
    3. Exploration: If no strong hypothesis guides you, choose an 'x' to explore a new region of the input space (e.g., negative numbers, non-integers, values between existing points) or to get more data points for pattern recognition.
    4. Rationale: Briefly explain *why* you chose this specific 'x' value, referencing the history, deductions, or hypotheses. State what you hope to learn from this query.

    - Always comment on the outcome of the *previous* query in your rationale (if applicable), relating it to your previous thoughts.
    - Use all available information to make an informed decision for the next query input.
    """

    what_you_know_thus_far = dspy.InputField(
        description="A structured summary (including confirmed properties, ruled-out properties, and hypotheses) of what you knew before this query."
    )
    query_history: list = dspy.InputField(
        description="List of previous queries and results (e.g., [{'input': 1, 'output': 3}, {'input': -1, 'output': -1}, ...])."
    )
    next_query_input: float = dspy.OutputField(
        description="The numerical input value 'x' you choose for the next query. Must be a valid number.",
    )


prober_agent = dspy.ChainOfThought(FunctionProber)


class FunctionKnowledgeUpdater(dspy.Signature):
    """
    You are a knowledge synthesizer analyzing data from function queries (input x, output f(x)).
    Your goal is to maintain and refine a list of deductions and hypotheses about the unknown function's properties based on the query history and the reasoning ('rationale') behind the last query.

    # Input:
    - rationale: The reasoning provided by the FunctionProber for choosing the *last* query input 'x'.
    - query_history: The *complete* list of {'input': x, 'output': f(x)} pairs gathered so far, including the latest result.
    - previous_thoughts: The list of confirmed properties, ruled-out properties, and hypotheses from the *previous* round.

    # Your Task: Update the 'thoughts' list (knowledge summary) for the *next* round.

    # Process:
    1. Integrate Latest Result: Analyze the newest (x, f(x)) pair in the context of the 'query_history'.
    2. Re-evaluate Existing Thoughts:
        * Carefully examine *every* existing "Confirmed property", "Ruled-out property", and "Hypothesis" from the 'previous_thoughts' against the *entire* 'query_history'.
        * **Confirmation Rule:** If a hypothesis is strongly supported by multiple (e.g., 3+) consistent data points and the overall pattern (e.g., points are collinear, differences suggest a specific polynomial degree, growth looks exponential), upgrade it to "Confirmed property: <summary>". State the evidence briefly (e.g., "Confirmed property: Linear relationship y=2x+3 holds for x=0, 1, 2").
        * **Refutation Rule:** If any data point in 'query_history' clearly contradicts a hypothesis or even a previous "Confirmed property", downgrade/change it to "Ruled-out property: <summary>". Explain the contradiction (e.g., "Ruled-out property: Function is not linear, as f(0)=1, f(1)=3, but f(2)=7"). Be ready to correct prior assumptions.
        * **Maintain Consistency:** If a property (confirmed or ruled-out) remains consistent with all data, keep it. If a hypothesis is neither strongly confirmed nor refuted by the current data, keep it as "Hypothesis: <summary>", possibly refining its description based on new data.

    3. Generate New Hypotheses:
        * Based on the overall shape suggested by the 'query_history' (plotting points mentally or analytically) and the prober's 'rationale', are there *new* potential properties or function types?
        * Consider: Linearity, specific polynomial degrees, exponential/logarithmic behavior, periodicity, intercepts (f(0)), roots (f(x)=0), symmetry (f(x)=f(-x) or f(x)=-f(-x)), asymptotes, domain/range limitations indicated by 'nan' or 'inf' outputs.
        * Formulate these as concise "Hypothesis: <new hypothesis>". Aim for testable hypotheses.

    4. Consolidate and Format:
        * Remove redundant or clearly invalidated thoughts.
        * If multiple hypotheses point to the same underlying idea, try to merge them.
        * Present the updated list clearly, grouping by "Confirmed property", "Ruled-out property", and "Hypothesis". If a category is empty, omit it.

    # Output Format:
    Return the updated knowledge summary as a bullet point list, like:
    * Confirmed property: Function passes through (0, 3).
    * Ruled-out property: Function is not quadratic (based on differences).
    * Hypothesis: Function might be linear (y=2x+3?).
    * Hypothesis: Function is increasing for positive x.
    """

    rationale: str = dspy.InputField(
        description="Rationale provided by the FunctionProber for their most recent query.",
    )
    query_history: list = dspy.InputField(
        description="Complete history of {'input': x, 'output': f(x)} queries and results.",
    )
    previous_thoughts: str = dspy.InputField(
        description="Bulleted list of confirmed properties, ruled-out properties, and hypotheses from the previous round.",
    )
    updated_thoughts: str = dspy.OutputField(
        description="Revised bullet point list of confirmed properties, ruled-out properties, and hypotheses for the next round.",
    )


knowledge_updater = dspy.Predict(FunctionKnowledgeUpdater)


# --- Game Loop ---
ROUNDS = 30  # Number of queries allowed

function_name, hidden_function = get_hidden_function(
    "x*sin(pi*x)"
)  # Change to desired function

# Initialize game state
thoughts = "* Hypothesis: The function exists and maps real numbers to real numbers (or NaN/inf)."  # Initial thought
query_history = []

print("--- Function Discovery Game ---")
print(
    f"Objective: Discover the properties of the hidden function f(x) in {ROUNDS} queries."
)
print("--------------------------------\n")


for i in range(ROUNDS):
    print(f"--- Round {i + 1}/{ROUNDS} ---")
    print(f"Current Knowledge:\n{thoughts if thoughts else 'None'}")
    # print(f"Query History: {query_history}") # Optional: Print full history each time

    # 1. Prober decides the next input 'x'
    probe_result = prober_agent(
        what_you_know_thus_far=thoughts,
        query_history=query_history,
    )

    # Attempt to parse the numerical input, handle potential errors
    try:
        # The LM might output text like "Query Input: 5.0". Extract the number.
        reasoning_text = probe_result.reasoning  # Get reasoning before modifying output
        # Try to find and convert the number in the output field
        # This part might need adjustment depending on the LM's exact output format
        query_input_str = str(probe_result.next_query_input).split(":")[-1].strip()
        next_x = float(query_input_str)

    except ValueError:
        print(
            "Error: Prober did not provide a valid numerical input. Choosing x=0 as fallback."
        )
        next_x = 0.0  # Fallback query
        reasoning_text = (
            probe_result.rationale + "\n[Error: Invalid input format, defaulted to x=0]"
        )

    print(f"\nProber's Reasoning:\n{reasoning_text}")
    print(f"Prober chose to query: x = {next_x}")

    # 2. Evaluate the hidden function
    try:
        output_y = hidden_function(next_x)
        # Handle potential non-finite outputs gracefully for history
        if math.isinf(output_y):
            output_y_str = "infinity" if output_y > 0 else "-infinity"
            output_y_for_history = float("inf") if output_y > 0 else float("-inf")
        elif math.isnan(output_y):
            output_y_str = "undefined (NaN)"
            output_y_for_history = float("nan")
        else:
            # Optional: Round for cleaner display/history if desired
            output_y = round(output_y, 4)
            output_y_str = str(output_y)
            output_y_for_history = output_y

        print(f"Function Result: f({next_x}) = {output_y_str}")

    except Exception as e:
        print(f"Error evaluating function at x={next_x}: {e}")
        output_y_str = "error"
        output_y_for_history = "error"  # Or perhaps nan? Decide how to represent this

    # 3. Update query history
    # Only add to history if it wasn't an evaluation error before calling f(x)
    if output_y_str != "error":
        current_query = {"input": next_x, "output": output_y_for_history}
        query_history.append(current_query)

    # 4. Update knowledge summary
    # Make sure to pass the *previous* thoughts state to the updater
    update_result = knowledge_updater(
        rationale=reasoning_text,
        query_history=query_history,  # Pass the new history
        previous_thoughts=thoughts,
    )
    thoughts = update_result.updated_thoughts

    print("\n--------------------------------\n")


# --- End of Game ---
print("--- Game Over ---")
print(f"Final Knowledge Summary:\n{thoughts}")
print(f"\nActual Hidden Function was: {function_name}")
print(f"Final Query History ({len(query_history)} points):")
for point in query_history:
    out_val = point["output"]
    if isinstance(out_val, float):
        if math.isinf(out_val):
            out_val = "inf" if out_val > 0 else "-inf"
        elif math.isnan(out_val):
            out_val = "nan"
        else:
            out_val = f"{out_val:.4f}"  # Format floats
    print(f"  Input: {point['input']}, Output: {out_val}")


# produce a plot of the function and the points

x_values = np.linspace(-10, 10, 400)
y_values = [hidden_function(x) for x in x_values]

query_x = [point["input"] for point in query_history]
query_y = [point["output"] for point in query_history]
plt.scatter(query_x, query_y, color="red", label="Query Points", zorder=5)
plt.plot(x_values, y_values, label=f"Hidden Function: {function_name}", zorder=1)
plt.title("Function Discovery Game")
plt.xlabel("Input (x)")
plt.ylabel("Output (f(x))")
plt.axhline(0, color="black", lw=0.5, ls="--")
plt.axvline(0, color="black", lw=0.5, ls="--")
plt.legend()
plt.grid()
plt.show()