#https://python.langchain.com/docs/how_to/functions/
import gradio as gr
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda

def word_counter(text):
    return len(text.split())

def process_with_function(text, api_key):
    try:
        # Initialize the model
        model = ChatOpenAI(
            openai_api_key=api_key,
            model="gpt-4o-mini"
        )
        
        # Create prompt template
        prompt = ChatPromptTemplate.from_template("Write a short story about {topic}")
        
        # Create chain with custom function
        chain = (
            prompt 
            | model 
            | RunnableLambda(lambda x: {
                "story": x.content,
                "word_count": word_counter(x.content)
            })
        )
        
        # Get result
        result = chain.invoke({"topic": text})
        
        return f"""
Story: {result['story']}

Word Count: {result['word_count']} words
"""
        
    except Exception as e:
        return f"Error: {str(e)}"

# Create Gradio interface
demo = gr.Interface(
    fn=process_with_function,
    inputs=[
        gr.Textbox(label="Story Topic", placeholder="Enter a topic..."),
        gr.Textbox(label="OpenAI API Key", type="password")
    ],
    outputs=gr.Textbox(label="Story with Analysis", lines=8),
    title="LangChain Functions Demo",
    description="Generates a story and counts its words using RunnableLambda"
)

if __name__ == "__main__":
    demo.launch()