import os
from orionagent import Agent, Manager, OpenAI, chat

# OpenRouter Configuration
OPENROUTER_API_KEY = "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"
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"
MODEL_NAME = "stepfun/step-3.5-flash:free"

def test_openrouter():
    print(f"Testing OpenAI Provider via OpenRouter with model: {MODEL_NAME}")
    
    # Initialize LLM
    llm = OpenAI(
        model_name=MODEL_NAME,
        api_key=OPENROUTER_API_KEY,
        base_url=OPENROUTER_BASE_URL,
        token_count=True,
        verbose=True,
        streaming=True # Re-enable streaming to verify fix
    )
    
    # Create a simple agent
    agent = Agent(
        name="Tester",
        role="Verification Expert",
        model=llm,
        system_instruction="You are a verification expert. Be concise."
    )
    
    # Run a simple chat
    print("\n--- Basic Chat Test ---")
    response = agent.ask("What is 1+1? Answer in one word.", stream=False)
    print(f"Response: {response}")
    
    # Run a multi-agent test with Manager
    print("\n--- Manager/Orchestration Test ---")
    manager = Manager(
        agents=[agent],
        model=llm,
        strategy="direct",
        verbose=True
    )
    
    mgr_response = manager.ask("Explain why 2+2=4 in one sentence.", stream=False)
    print(f"Manager Response: {mgr_response}")
    
    llm.print_session_tokens()

if __name__ == "__main__":
    test_openrouter()
