import sdialog
from sdialog import Context
from sdialog.agents import Agent
from sdialog.personas import Persona
from sdialog.orchestrators import SimpleReflexOrchestrator

# First, let's set our preferred backend/model and parameters
sdialog.config.llm("openai:gpt-4.1", temperature=0.9)

# Let's define our personas
alice = Persona(name="Alice", role="barista", personality="cheerful")
bob   = Persona(name="Bob", role="customer", personality="curious")

# (Optional) Let's add a concrete conversational context
ctx = Context(
  location="Downtown cafe",
  environment="noisy, aromatic cafe with occasional grinder sounds",
  circumstances="Morning rush hour",
  objects=["espresso machine", "menu board", "tip jar"]
)

# (Optional) Let's add a simple tool (just a plain Python function)
# We'll use a tiny mock function our agent can call as a tool
def lookup_menu(item: str) -> dict:
    return {"item": item, "specials": ["vanilla latte", "cold brew"]}

# (Optional) Let's include a small rule-based orchestrator
react = SimpleReflexOrchestrator(
    condition=lambda utt: "decaf" in utt.lower(),
    instruction="Explain decaf options and suggest one."
)

# Now we create the agents
barista = Agent(persona=alice, tools=[lookup_menu])
customer = Agent(persona=bob, first_utterance="Hi!")

# (Optional) We can attach orchestrators to an agent using pipe-like composition
barista = barista | react

# Let's generate three dialogs!
for ix in range(3):
    dialog = customer.dialog_with(barista, context=ctx)
    dialog.print(orchestration=True)
    dialog.to_file(f"dialog_{ix}.json")

from sdialog import Dialog

# Load from JSON (generated by SDialog using `to_file()`)
dialog = Dialog.from_file("dialog_0.json")

# Load from HuggingFace Hub datasets
dialogs = Dialog.from_huggingface("sdialog/Primock-57")

# Create from plain text files or strings - perfect for converting existing datasets!
dialog_from_txt = Dialog.from_str("""
Alice: Hello there! How are you today?
Bob: I'm doing great, thanks for asking.
Alice: That's wonderful to hear!
""")
# Or, equivalently if the content is in a txt file
dialog_from_txt = Dialog.from_file("conversation.txt")

# Load from CSV files with custom templates
dialog_from_csv = Dialog.from_file("conversation.csv",
                                   csv_speaker_col="speaker",
                                   csv_text_col="value",)

# All Dialog objects have rich manipulation methods
dialog.filter("Alice").rename_speaker("Alice", "Customer").upper().to_file("processed.json")
avg_words_turn = sum(len(turn) for turn in dialog) / len(dialog)

from sdialog.personas import Doctor, Patient
from sdialog.generators import PersonaGenerator, ContextGenerator
from sdialog import Context

# By default, unspecified attributes are LLM generated
doc = PersonaGenerator(Doctor(specialty="Cardiology")).generate()
pat = PersonaGenerator(Patient(symptoms="chest pain")).generate()

# Optionally specify generation sources per attribute
ctx_gen = ContextGenerator(Context(location="emergency room"))
ctx_gen.set(
  objects=get_random_object,                        # user-defined function
  circumstances="{csv:circumstance:./data/circumstances.csv}",  # CSV file values
  goals="{llm:Suggest a realistic goal for the context}"         # targeted LLM instruction
)
ctx = ctx_gen.generate()

# Our remote system (your conversational backend exposing an OpenAI-compatible API)
system = Agent(
  model="my/super-llm",  # Model name exposed by your server
  openai_api_base="http://my-endpoint.com:8000/v1",  # Base URL of the service
  openai_api_key="EMPTY",  # Or a real key if required
  name="System"
)

# Let's manually define one (minimal) synthetic customer persona
angry_customer = Customer(
  name="Riley",
  issue="Billing error on last invoice",
  issue_description="Charged twice for the same month",
  anger_level="high",
  times_called=3,
)

simulated_customer = Agent(persona=angry_customer, name="Customer")

# Let's make the system talk to our simulated customer once
dialog = system.dialog_with(simulated_customer)
dialog.to_file("dialog_0.json")

from sdialog.evaluation import LLMJudgeRealDialog, LinguisticFeatureScore
from sdialog.evaluation import FrequencyEvaluator, MeanEvaluator
from sdialog.evaluation import DatasetComparator

reference = [...]   # list[Dialog]
candidate = [...]   # list[Dialog]

judge  = LLMJudgeRealDialog()
flesch = LinguisticFeatureScore(feature="flesch-reading-ease")

comparator = DatasetComparator([
  FrequencyEvaluator(judge, name="Realistic dialog rate"),
  MeanEvaluator(flesch, name="Mean Flesch Reading Ease"),
])

results = comparator({"reference": reference, "candidate": candidate})

# Plot results for each evaluator
comparator.plot()

import sdialog
from sdialog.interpretability import Inspector
from sdialog.agents import Agent

sdialog.config.llm("huggingface:meta-llama/Llama-3.2-3B-Instruct")

agent = Agent(name="Bob")
inspector = Inspector(target="model.layers.16.post_attention_layernorm")
agent = agent | inspector

agent("How are you?")
agent("Cool!")

# Let's get the last response's first token activation vector!
act = inspector[-1][0].act # [response index][token index]

anger_direction = torch.load("anger_direction.pt")  # A direction vector (e.g., PCA / difference-in-mean vector)
agent_steered = agent | inspector - anger_direction  # Ablate the anger direction from the target activations

agent_steered("You are an extremely upset assistant")  # Agent "can't get angry anymore" :)

import sdialog

# Change the default global LLM
sdialog.config.llm("ollama:qwen3:14b")
# Any argument supported by the chosen backend/model can also be given, for example
sdialog.config.llm("ollama:qwen3:14b",
                   temperature=0.7,
                   base_url="https://my-ollama-endpoint.com:123")  # Remote Ollama server

from sdialog.agents import Agent

my_agent = Agent(model="amazon:anthropic.claude-3-5-sonnet-20240620-v1:0",
                 region_name="us-east-1")