Update app.py

app.py

@@ -1,33 +1,108 @@
 import gradio as gr
 from gradio import ChatMessage
-from transformers import load_tool, ReactCodeAgent, HfEngine
-from utils import stream_from_transformers_agent
+from typing import Iterator
+import google.generativeai as genai
+import os 
 
-# Import tool from Hub
-image_generation_tool = load_tool("m-ric/text-to-image")
+genai.configure(api_key=os.environ.get("GEMINI_API_KEY"))
 
+model = genai.GenerativeModel("gemini-2.0-flash-thinking-exp-1219")
 
-llm_engine = HfEngine("meta-llama/Meta-Llama-3-70B-Instruct")
-# Initialize the agent with both tools
-agent = ReactCodeAgent(tools=[image_generation_tool], llm_engine=llm_engine)
-
-
-def interact_with_agent(prompt, messages):
-    messages.append(ChatMessage(role="user", content=prompt))
-    yield messages
-    for msg in stream_from_transformers_agent(agent, prompt):
-        messages.append(msg)
+def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
+    """
+    Streams both thoughts and responses from the Gemini model.
+    """
+    # Initialize response from Gemini
+    response = model.generate_content(user_message, stream=True)
+    
+    # Initialize buffers
+    thought_buffer = ""
+    response_buffer = ""
+    thinking_complete = False
+    
+    # Add initial thinking message
+    messages.append(
+        ChatMessage(
+            role="assistant",
+            content="",
+            metadata={"title": "⏳Thinking: *The thoughts produced by the Gemini2.0 Flash model are experimental"}
+        )
+    )
+    
+    for chunk in response:
+        parts = chunk.candidates[0].content.parts
+        current_chunk = parts[0].text
+        
+        if len(parts) == 2 and not thinking_complete:
+            # Complete thought and start response
+            thought_buffer += current_chunk
+            messages[-1] = ChatMessage(
+                role="assistant",
+                content=thought_buffer,
+                metadata={"title": "⏳Thinking: *The thoughts produced by the Gemini2.0 Flash model are experimental"}
+            )
+            
+            # Add response message
+            messages.append(
+                ChatMessage(
+                    role="assistant",
+                    content=parts[1].text
+                )
+            )
+            thinking_complete = True
+            
+        elif thinking_complete:
+            # Continue streaming response
+            response_buffer += current_chunk
+            messages[-1] = ChatMessage(
+                role="assistant",
+                content=response_buffer
+            )
+            
+        else:
+            # Continue streaming thoughts
+            thought_buffer += current_chunk
+            messages[-1] = ChatMessage(
+                role="assistant",
+                content=thought_buffer,
+                metadata={"title": "⏳Thinking: *The thoughts produced by the Gemini2.0 Flash model are experimental"}
+            )
+        
         yield messages
-    yield messages
 
 
 with gr.Blocks() as demo:
-    chatbot = gr.Chatbot(label="Agent",
-                         msg_format="messages",
-                         avatar_images=(None, "https://em-content.zobj.net/source/twitter/53/robot-face_1f916.png"))
-    text_input = gr.Textbox(lines=1, label="Chat Message")
-    text_input.submit(interact_with_agent, [text_input, chatbot], [chatbot])
-
+    gr.Markdown("# Chat with Gemini 2.0 Flash and See its Thoughts 💭")
+    
+    chatbot = gr.Chatbot(
+        type="messages",
+        label="Gemini2.0 'Thinking' Chatbot",
+        render_markdown=True,
+    )
+    
+    input_box = gr.Textbox(
+        lines=1,
+        label="Chat Message",
+        placeholder="Type your message here and press Enter..."
+    )
+    
+    # Set up event handlers
+    msg_store = gr.State("")  # Store for preserving user message
+    
+    input_box.submit(
+        lambda msg: (msg, msg, ""),  # Store message and clear input
+        inputs=[input_box],
+        outputs=[msg_store, input_box, input_box],
+        queue=False
+    ).then(
+        user_message,  # Add user message to chat
+        inputs=[msg_store, chatbot],
+        outputs=[input_box, chatbot],
+        queue=False
+    ).then(
+        stream_gemini_response,  # Generate and stream response
+        inputs=[msg_store, chatbot],
+        outputs=chatbot
+    )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()