ysharma HF staff commited on
Commit
4a3acd2
·
verified ·
1 Parent(s): 772fb5d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +41 -27
app.py CHANGED
@@ -12,91 +12,104 @@ genai.configure(api_key=GEMINI_API_KEY)
12
  model = genai.GenerativeModel("gemini-2.0-flash-thinking-exp-1219")
13
 
14
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
15
  def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
16
  """
17
- Streams LLM's thoughts and response.
18
  """
19
  try:
20
- # Enabling logging for users to understand how thinking works along with streaming
21
  print(f"\n=== New Request ===")
22
  print(f"User message: {user_message}")
23
-
24
- # Initialize response from Gemini
25
- response = model.generate_content(user_message, stream=True)
26
-
 
 
 
 
27
  # Initialize buffers and flags
28
  thought_buffer = ""
29
  response_buffer = ""
30
- has_response = False
31
  thinking_complete = False
32
-
33
  # Add initial thinking message
34
  messages.append(
35
  ChatMessage(
36
  role="assistant",
37
  content="",
38
- metadata={"title": "Thinking: *The thoughts produced by the model are experimental"}
39
  )
40
  )
41
-
42
  for chunk in response:
43
  parts = chunk.candidates[0].content.parts
44
  current_chunk = parts[0].text
45
-
46
  if len(parts) == 2 and not thinking_complete:
47
- # Complete thought
48
  thought_buffer += current_chunk
49
  print(f"\n=== Complete Thought ===\n{thought_buffer}")
50
-
51
- # Update thinking message
52
  messages[-1] = ChatMessage(
53
  role="assistant",
54
  content=thought_buffer,
55
- metadata={"title": "Thinking: *The thoughts produced by the model are experimental"}
56
  )
57
  yield messages
 
58
 
59
  # Start response
60
  response_buffer = parts[1].text
61
  print(f"\n=== Starting Response ===\n{response_buffer}")
62
-
63
  messages.append(
64
  ChatMessage(
65
  role="assistant",
66
  content=response_buffer
67
  )
68
  )
69
-
70
  thinking_complete = True
71
- has_response = True
72
- yield messages
73
 
74
  elif thinking_complete:
75
  # Stream response
76
  response_buffer += current_chunk
77
  print(f"\n=== Response Chunk ===\n{current_chunk}")
78
-
79
  messages[-1] = ChatMessage(
80
  role="assistant",
81
  content=response_buffer
82
  )
83
- yield messages
84
 
85
  else:
86
  # Stream thinking
87
  thought_buffer += current_chunk
88
  print(f"\n=== Thinking Chunk ===\n{current_chunk}")
89
-
90
  messages[-1] = ChatMessage(
91
  role="assistant",
92
  content=thought_buffer,
93
- metadata={"title": "Thinking: *The thoughts produced by the model are experimental"}
94
  )
95
- yield messages
96
-
97
- # Log final complete response
 
98
  print(f"\n=== Final Response ===\n{response_buffer}")
99
-
100
  except Exception as e:
101
  print(f"\n=== Error ===\n{str(e)}")
102
  messages.append(
@@ -111,6 +124,7 @@ def user_message(msg: str, history: list) -> tuple[str, list]:
111
  """Adds user message to chat history"""
112
  history.append(ChatMessage(role="user", content=msg))
113
  return "", history
 
114
 
115
  # Create the Gradio interface
116
  with gr.Blocks(theme=gr.themes.Citrus(), fill_height=True) as demo:
 
12
  model = genai.GenerativeModel("gemini-2.0-flash-thinking-exp-1219")
13
 
14
 
15
+ def format_chat_history(messages: list) -> list:
16
+ """
17
+ Formats the chat history into a structure Gemini can understand
18
+ """
19
+ formatted_history = []
20
+ for message in messages:
21
+ # Skip thinking messages (messages with metadata)
22
+ if not (message.get("role") == "assistant" and "metadata" in message):
23
+ formatted_history.append({
24
+ "role": "user" if message.get("role") == "user" else "assistant",
25
+ "parts": [message.get("content", "")]
26
+ })
27
+ return formatted_history
28
+
29
  def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
30
  """
31
+ Streams thoughts and response with conversation history support.
32
  """
33
  try:
 
34
  print(f"\n=== New Request ===")
35
  print(f"User message: {user_message}")
36
+
37
+ # Format chat history for Gemini
38
+ chat_history = format_chat_history(messages)
39
+
40
+ # Initialize Gemini chat
41
+ chat = model.start_chat(history=chat_history)
42
+ response = chat.send_message(user_message, stream=True)
43
+
44
  # Initialize buffers and flags
45
  thought_buffer = ""
46
  response_buffer = ""
 
47
  thinking_complete = False
48
+
49
  # Add initial thinking message
50
  messages.append(
51
  ChatMessage(
52
  role="assistant",
53
  content="",
54
+ metadata={"title": "⚙️ Thinking: *The thoughts produced by the model are experimental"}
55
  )
56
  )
57
+
58
  for chunk in response:
59
  parts = chunk.candidates[0].content.parts
60
  current_chunk = parts[0].text
61
+
62
  if len(parts) == 2 and not thinking_complete:
63
+ # Complete thought and start response
64
  thought_buffer += current_chunk
65
  print(f"\n=== Complete Thought ===\n{thought_buffer}")
66
+
 
67
  messages[-1] = ChatMessage(
68
  role="assistant",
69
  content=thought_buffer,
70
+ metadata={"title": "⚙️ Thinking: *The thoughts produced by the model are experimental"}
71
  )
72
  yield messages
73
+ time.sleep(0.05)
74
 
75
  # Start response
76
  response_buffer = parts[1].text
77
  print(f"\n=== Starting Response ===\n{response_buffer}")
78
+
79
  messages.append(
80
  ChatMessage(
81
  role="assistant",
82
  content=response_buffer
83
  )
84
  )
 
85
  thinking_complete = True
 
 
86
 
87
  elif thinking_complete:
88
  # Stream response
89
  response_buffer += current_chunk
90
  print(f"\n=== Response Chunk ===\n{current_chunk}")
91
+
92
  messages[-1] = ChatMessage(
93
  role="assistant",
94
  content=response_buffer
95
  )
 
96
 
97
  else:
98
  # Stream thinking
99
  thought_buffer += current_chunk
100
  print(f"\n=== Thinking Chunk ===\n{current_chunk}")
101
+
102
  messages[-1] = ChatMessage(
103
  role="assistant",
104
  content=thought_buffer,
105
+ metadata={"title": "⚙️ Thinking: *The thoughts produced by the model are experimental"}
106
  )
107
+
108
+ yield messages
109
+ time.sleep(0.05)
110
+
111
  print(f"\n=== Final Response ===\n{response_buffer}")
112
+
113
  except Exception as e:
114
  print(f"\n=== Error ===\n{str(e)}")
115
  messages.append(
 
124
  """Adds user message to chat history"""
125
  history.append(ChatMessage(role="user", content=msg))
126
  return "", history
127
+
128
 
129
  # Create the Gradio interface
130
  with gr.Blocks(theme=gr.themes.Citrus(), fill_height=True) as demo: