Psiska commited on
Commit
387433f
Β·
1 Parent(s): 5d85d64

commented code no Conversational Memory

Browse files
Files changed (2) hide show
  1. agent.py +94 -22
  2. app.py +14 -1
agent.py CHANGED
@@ -6,52 +6,97 @@ import tools.tools as tls # Your tool definitions
6
 
7
  load_dotenv()
8
 
9
- # βœ… Utility: Ensure role alternation (user/assistant/user...)
10
- def enforce_strict_role_alternation(messages):
11
- """
12
- Fixes message history to enforce 'user/assistant/user/assistant' alternation.
13
- Keeps first 'system' if present.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
14
  """
15
- cleaned = []
16
- last_role = None
 
17
 
18
  for msg in messages:
19
  role = msg["role"]
 
 
20
  if role not in ("user", "assistant", "system"):
21
- continue # skip invalid roles
22
 
 
23
  if role == "system" and not cleaned:
24
  cleaned.append(msg)
25
  continue
26
 
 
27
  if role == last_role:
28
- continue # skip consecutive same-role messages
29
 
 
30
  cleaned.append(msg)
31
- last_role = role
32
 
33
  return cleaned
34
 
35
 
36
- # βœ… Custom Model Wrapper using InferenceClient
37
  class HuggingFaceChatModel(Model):
38
  def __init__(self):
 
39
  model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
 
 
40
  self.client = InferenceClient(model=model_id, token=os.getenv("HF_TOKEN"))
41
 
42
  def generate(self, messages, stop_sequences=None):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
43
  if stop_sequences is None:
44
  stop_sequences = ["Task"]
45
 
46
- # πŸ’‘ Enforce correct message order
47
  cleaned_messages = enforce_strict_role_alternation(messages)
48
 
49
- # πŸ”§ Hugging Face call
50
  response = self.client.chat_completion(
51
  messages=cleaned_messages,
52
  stop=stop_sequences,
53
- max_tokens=1024
54
  )
 
 
55
  content = response.choices[0].message["content"]
56
  return ChatMessage(role="assistant", content=content)
57
 
@@ -59,25 +104,52 @@ class HuggingFaceChatModel(Model):
59
  # βœ… Basic Agent with SmolAgents
60
  class BasicAgent:
61
  def __init__(self):
 
62
  print("βœ… BasicAgent initialized with Hugging Face chat model.")
 
 
63
  self.model = HuggingFaceChatModel()
64
 
 
65
  self.agent = CodeAgent(
66
- tools=[tls.search_tool, tls.calculate_cargo_travel_time],
67
- model=self.model,
68
- additional_authorized_imports=["pandas"],
69
- max_steps=20,
70
  )
71
 
72
  def __call__(self, messages) -> str:
73
- # Extract last user message from the chat history
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
74
  if isinstance(messages, list):
75
- question = messages[-1]["content"] # Gradio gives list of dicts
76
  else:
77
- question = messages # Fallback if it's just a string
78
 
 
79
  print(f"πŸ“₯ Received question: {question[:60]}...")
 
 
80
  response = self.agent.run(question)
 
 
81
  print(f"πŸ“€ Response generated: {response[:60]}...")
82
- return response
 
 
83
 
 
6
 
7
  load_dotenv()
8
 
9
+ """
10
+ enforce_strict_role_alternation()
11
+
12
+ Ensures that messages follow the required pattern:
13
+ 'user/assistant/user/assistant/...', starting with an optional 'system' message.
14
+
15
+ This is necessary because many chat-based models (e.g., ChatCompletion APIs)
16
+ expect the conversation format to alternate strictly between user and assistant roles,
17
+ possibly preceded by a single system message.
18
+
19
+ Parameters:
20
+ -----------
21
+ messages : list of dict
22
+ The message history. Each message is expected to be a dictionary with a 'role' key
23
+ ('user', 'assistant', or 'system') and a 'content' key.
24
+
25
+ Returns:
26
+ --------
27
+ cleaned : list of dict
28
+ A sanitized version of the messages list that follows the correct role alternation rules.
29
  """
30
+ def enforce_strict_role_alternation(messages):
31
+ cleaned = [] # List to store the cleaned message sequence
32
+ last_role = None # Tracks the last valid role added to ensure alternation
33
 
34
  for msg in messages:
35
  role = msg["role"]
36
+
37
+ # Skip any message that doesn't have a valid role
38
  if role not in ("user", "assistant", "system"):
39
+ continue
40
 
41
+ # Allow a single 'system' message only at the very beginning
42
  if role == "system" and not cleaned:
43
  cleaned.append(msg)
44
  continue
45
 
46
+ # Skip messages with the same role as the previous one (breaks alternation)
47
  if role == last_role:
48
+ continue
49
 
50
+ # Add the valid message to the cleaned list
51
  cleaned.append(msg)
52
+ last_role = role # Update the last role for the next iteration
53
 
54
  return cleaned
55
 
56
 
57
+ # Define a custom model class that wraps around Hugging Face's InferenceClient for chat-based models
58
  class HuggingFaceChatModel(Model):
59
  def __init__(self):
60
+ # Set the model ID for the specific Hugging Face model to use
61
  model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
62
+
63
+ # Create an InferenceClient with the model ID and the Hugging Face token from your environment
64
  self.client = InferenceClient(model=model_id, token=os.getenv("HF_TOKEN"))
65
 
66
  def generate(self, messages, stop_sequences=None):
67
+ """
68
+ Generates a response from the chat model based on the input message history.
69
+
70
+ Parameters:
71
+ -----------
72
+ messages : list of dict
73
+ A list of message dicts in OpenAI-style format, e.g.:
74
+ [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi!"}]
75
+
76
+ stop_sequences : list of str, optional
77
+ A list of strings that will stop generation when encountered. Default is ["Task"].
78
+
79
+ Returns:
80
+ --------
81
+ ChatMessage
82
+ A formatted response object with role='assistant' and the model-generated content.
83
+ """
84
+
85
+ # Set default stop sequences if none provided
86
  if stop_sequences is None:
87
  stop_sequences = ["Task"]
88
 
89
+ # πŸ’‘ Preprocess: Enforce valid alternation of user/assistant messages
90
  cleaned_messages = enforce_strict_role_alternation(messages)
91
 
92
+ # πŸ”§ Call the Hugging Face chat API with cleaned messages
93
  response = self.client.chat_completion(
94
  messages=cleaned_messages,
95
  stop=stop_sequences,
96
+ max_tokens=1024 # Limit the number of tokens generated in the reply
97
  )
98
+
99
+ # πŸ“¦ Extract content from the model response and wrap it in a ChatMessage object
100
  content = response.choices[0].message["content"]
101
  return ChatMessage(role="assistant", content=content)
102
 
 
104
  # βœ… Basic Agent with SmolAgents
105
  class BasicAgent:
106
  def __init__(self):
107
+ # Informative log to indicate that the agent is being initialized
108
  print("βœ… BasicAgent initialized with Hugging Face chat model.")
109
+
110
+ # Instantiate your custom model that wraps the Hugging Face InferenceClient
111
  self.model = HuggingFaceChatModel()
112
 
113
+ # Create the CodeAgent, which uses the tools and the chat model
114
  self.agent = CodeAgent(
115
+ tools=[tls.search_tool, tls.calculate_cargo_travel_time], # Your list of tools
116
+ model=self.model, # The model to generate tool-using responses
117
+ additional_authorized_imports=["pandas"], # Optional: allow use of pandas in generated code
118
+ max_steps=20, # Limit the number of planning steps (tool calls + reasoning)
119
  )
120
 
121
  def __call__(self, messages) -> str:
122
+ """
123
+ Handle a call to the agent with either a single question or a message history.
124
+
125
+ Parameters:
126
+ -----------
127
+ messages : Union[str, List[Dict[str, str]]]
128
+ The input from the chat interface β€” either:
129
+ - a plain string (just one message)
130
+ - a list of dicts, like [{"role": "user", "content": "What's the weather?"}]
131
+
132
+ Returns:
133
+ --------
134
+ str
135
+ The assistant's response as a string.
136
+ """
137
+
138
+ # If the input is a chat history (list of messages), get the most recent user message
139
  if isinstance(messages, list):
140
+ question = messages[-1]["content"] # Extract last message content
141
  else:
142
+ question = messages # If it's just a string, use it directly
143
 
144
+ # Log the input for debugging
145
  print(f"πŸ“₯ Received question: {question[:60]}...")
146
+
147
+ # Run the CodeAgent to get a response (may include tool use)
148
  response = self.agent.run(question)
149
+
150
+ # Log the response for debugging
151
  print(f"πŸ“€ Response generated: {response[:60]}...")
152
+
153
+ return response # Return final result
154
+
155
 
app.py CHANGED
@@ -147,18 +147,31 @@ def test_init_agent_for_chat(text_input, history):
147
 
148
  return submitted_answer
149
  '''
 
 
 
150
  def test_init_agent_for_chat(text_input, history):
151
  try:
 
 
152
  basicAgent = BasicAgent()
153
  except Exception as e:
 
154
  return f"[Error initializing agent]: {e}"
155
 
 
156
  return basicAgent(text_input)
157
 
 
158
  with gr.Blocks() as demo:
 
159
  gr.Markdown("## πŸ€– Conversational Cargo Agent")
 
 
 
160
  gr.ChatInterface(test_init_agent_for_chat, type="messages")
161
 
 
162
  # --- Build Gradio Interface using Blocks ---
163
  with gr.Blocks() as demo:
164
  gr.Markdown("# Basic Agent Evaluation Runner")
@@ -191,7 +204,7 @@ with gr.Blocks() as demo:
191
  # fn=run_and_submit_all,
192
  # outputs=[status_output, results_table]
193
  # )
194
-
195
 
196
  if __name__ == "__main__":
197
  load_dotenv()
 
147
 
148
  return submitted_answer
149
  '''
150
+ # βœ… This function is the core callback for the Gradio chat interface.
151
+ # It is called every time the user submits a new message.
152
+
153
  def test_init_agent_for_chat(text_input, history):
154
  try:
155
+ # πŸ”§ Try to initialize an instance of your BasicAgent.
156
+ # You could later refactor this to reuse a single instance instead of re-creating it every time.
157
  basicAgent = BasicAgent()
158
  except Exception as e:
159
+ # ❌ If initialization fails, return the error message to the user.
160
  return f"[Error initializing agent]: {e}"
161
 
162
+ # πŸ’¬ Pass the user input (text_input) to the agent and return the agent's response.
163
  return basicAgent(text_input)
164
 
165
+ # βœ… Define the Gradio app UI using Blocks (layout container).
166
  with gr.Blocks() as demo:
167
+ # πŸ“ Add a markdown title to the UI
168
  gr.Markdown("## πŸ€– Conversational Cargo Agent")
169
+
170
+ # πŸ’¬ Create a chat interface connected to the agent function.
171
+ # type="messages" ensures it receives and sends message history in OpenAI-style format.
172
  gr.ChatInterface(test_init_agent_for_chat, type="messages")
173
 
174
+ '''
175
  # --- Build Gradio Interface using Blocks ---
176
  with gr.Blocks() as demo:
177
  gr.Markdown("# Basic Agent Evaluation Runner")
 
204
  # fn=run_and_submit_all,
205
  # outputs=[status_output, results_table]
206
  # )
207
+ '''
208
 
209
  if __name__ == "__main__":
210
  load_dotenv()