# ============================================================================
#  Codex-as-API  ->  OpenAI Agents SDK  (personal "self AI" agent)
#  HOW TO USE: open https://colab.research.google.com  ->  new notebook  ->
#  paste this WHOLE file into ONE cell  ->  press Run.  That's it.
#
#  It does 4 things:
#    1. installs the OpenAI Agents SDK
#    2. points the SDK at YOUR Codex API (running on your Hugging Face Space)
#    3. runs an agent and proves session memory works
#    4. leaves you a reusable  ask("...")  helper for the next cells
# ============================================================================

# --- 1) install (quiet) ----------------------------------------------------
!pip -q install openai-agents nest_asyncio

import asyncio, nest_asyncio
from openai import AsyncOpenAI
from agents import Agent, Runner, OpenAIChatCompletionsModel, set_tracing_disabled

nest_asyncio.apply()          # lets asyncio.run() work inside Colab's loop

# --- 2) YOUR API settings (edit these) -------------------------------------
BASE_URL = "https://sarveshpatel-codex.hf.space/v1"   # your Space + /v1
API_KEY  = "CURSEOFWITCHER"                            # your API_TOKEN secret
MODEL    = "codex"                                     # always "codex"
SESSION  = "colab-self-ai"   # any string. Same string = persistent memory.

# --- 3) wire the Agents SDK to your Codex API ------------------------------
set_tracing_disabled(True)   # don't send traces to OpenAI

client = AsyncOpenAI(
    base_url=BASE_URL,
    api_key=API_KEY,
    default_headers={"X-Session-Id": SESSION},   # gives the agent memory
)
model = OpenAIChatCompletionsModel(model=MODEL, openai_client=client)

# This is your agent. Change `instructions` to define its personality/role.
self_ai = Agent(
    name="SelfAI",
    instructions=(
        "You are SelfAI, my personal coding and research assistant powered by "
        "Codex. Be concise, direct, and practical."
    ),
    model=model,
)

# --- 4) a reusable helper you can call in any later cell -------------------
def ask(prompt: str) -> str:
    """Send a message to your agent and return its reply (remembers the session)."""
    result = asyncio.run(Runner.run(self_ai, prompt))
    return result.final_output


# --- 5) quick demo: connectivity + memory ----------------------------------
print("ping  :", asyncio.run(
    client.chat.completions.create(
        model=MODEL, messages=[{"role": "user", "content": "Reply with exactly: PONG"}]
    )
).choices[0].message.content)

print("agent :", ask("In one sentence, what can you help me with?"))
print("memory>", ask("My favorite number is 42. Acknowledge in 3 words."))
print("recall:", ask("What is my favorite number? Reply with just the number."))

# ---------------------------------------------------------------------------
# NEXT CELLS — just call ask():
#     ask("Write a Python function to check if a string is a palindrome.")
#     ask("Now add a test for it.")          # it remembers the previous answer
#
# TIPS:
#   * New conversation? change SESSION (or set it to None for one-off calls).
#   * Streaming / tools: this API is OpenAI-compatible, so any OpenAI client
#     works. Note: OpenAI-style function tools aren't forwarded yet — but
#     Codex's own tools (running code, editing files) execute server-side
#     inside the session.
# ---------------------------------------------------------------------------