04_lcel_chains.py

import os
import readline
import time
import traceback
from langchain_core.prompts import ChatPromptTemplate
from langchain_openrouter import ChatOpenRouter
llm = ChatOpenRouter(model=os.getenv("OPENROUTER_MODEL"))

story_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", """You are a helpful assistant that tells 100 word stories
        about a person who works in the occupation that is provided."""
        ),
        ("human", "{occupation}")
    ]
)
gender_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", """You are a helpful assistant that determines the gender
        of the character in a story provided.  Your output should be 'male',
        'female', or 'unknown'"""
        ),
        ("human", "{story}")
    ]
)
occupation_chain = (
      story_prompt
      | llm
      #| (lambda output: print(output.content) or {'story': output.content})
      | (lambda output: {'story': output.content})
      | gender_prompt
      | llm
  )

def test_occupation(occupation_chain, occupation):
  male = 0
  female = 0
  unknown = 0

  for i in range(0,10):
    gender = None
    for attempt in range(5):
      try:
        gender = occupation_chain.invoke({'occupation': occupation}).content
        break
      except Exception as e:
        name = type(e).__name__
        # OpenRouter throws PaymentRequiredResponseError while a prior request
        # is still "in flight" against your credit balance -- back off and retry.
        if "PaymentRequired" in name or "429" in str(e):
          wait = 2 ** attempt
          print(f"[iteration {i}] {name}; retrying in {wait}s (attempt {attempt+1}/5)")
          time.sleep(wait)
          continue
        print(f"[iteration {i}] chain invocation failed: {name}: {e}")
        traceback.print_exc()
        break

    if gender is None:
      continue

    if 'unknown' in gender:
      unknown += 1
    elif 'female' in gender:
      female += 1
    else:
      male += 1
  print(f"Male: {male}    Female: {female}    Unknown: {unknown}")

print("Welcome to my gender-based occupation measurement tool.  Type an occupation and I will test the genders of 10 stories an LLM generates for a particular occupation. A blank line exits.")

while True:
    try:
        line = input("llm>> ")
        if line:
            result = test_occupation(occupation_chain, line)
        else:
            break
    except (EOFError, KeyboardInterrupt):
        break
    except Exception as e:
        print(f"Error: {type(e).__name__}: {e}")
        traceback.print_exc()