Update app.py

app.py

@@ -1,24 +1,93 @@
-from langchain_community.callbacks import get_openai_callback
+import gradio as gr
 from langchain_openai import ChatOpenAI
+from langchain_community.callbacks import get_openai_callback
+import os
 
-def track_tokens_and_cost(message):
-    llm = ChatOpenAI(temperature=0.7)
-    
-    with get_openai_callback() as cb:
-        response = llm.invoke(message)
+def chat_with_tracking(message, api_key):
+    try:
+        # Set API key
+        os.environ["OPENAI_API_KEY"] = api_key
         
-        # Access the correct attributes
-        print(f"""
-Response: {response.content}
+        # Initialize ChatOpenAI with error handling
+        try:
+            llm = ChatOpenAI(
+                model="gpt-3.5-turbo",  # Using 3.5-turbo as it's more reliable
+                temperature=0.7
+            )
+        except Exception as e:
+            return f"Error initializing model: {str(e)}"
+
+        # Track token usage with error handling
+        try:
+            with get_openai_callback() as cb:
+                response = llm.invoke(message)
+                
+                # Format the response with token usage
+                result = f"""
+🤖 Response: 
+{response.content}
 
-Token Usage:
------------
-Prompt Tokens: {cb.prompt_tokens}
-Completion Tokens: {cb.completion_tokens}
-Total Tokens: {cb.total_tokens}
+📊 Token Usage Statistics:
+-------------------------
+📥 Input Tokens: {cb.prompt_tokens}
+📤 Output Tokens: {cb.completion_tokens}
+📑 Total Tokens: {cb.total_tokens}
 
-Cost Analysis:
--------------
+💰 Cost Analysis:
+---------------
 Total Cost: ${cb.total_cost:.6f}
-        """)
-        return response
+"""
+                return result
+                
+        except Exception as e:
+            return f"Error during token tracking: {str(e)}"
+            
+    except Exception as e:
+        return f"General error: {str(e)}"
+    finally:
+        # Clear API key from environment for security
+        if "OPENAI_API_KEY" in os.environ:
+            del os.environ["OPENAI_API_KEY"]
+
+# Create Gradio interface with better styling
+demo = gr.Interface(
+    fn=chat_with_tracking,
+    inputs=[
+        gr.Textbox(
+            label="Your Message",
+            placeholder="Type your message here...",
+            lines=4,
+            scale=2
+        ),
+        gr.Textbox(
+            label="OpenAI API Key",
+            placeholder="Enter your OpenAI API key",
+            type="password",
+            scale=1
+        )
+    ],
+    outputs=gr.Textbox(
+        label="Response and Analytics",
+        lines=10
+    ),
+    title="🤖 LangChain Token Usage Tracker",
+    description="Track token usage and costs for your conversations with GPT",
+    theme="default",
+    css="""
+        .gradio-container {
+            font-family: 'Arial', sans-serif;
+        }
+        .output-box {
+            font-family: 'Courier New', monospace;
+        }
+    """
+)
+
+# Launch the application
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        debug=True
+    )