blackopsrepl commited on
Commit
e2685f7
Β·
1 Parent(s): f3473c1

feat!: add chatbot demo ui

Browse files

- Add MCP Chat Interface Usage section with natural language examples
- Enhance Architecture section with MCPClientWrapper and MCP components
- Add 5 new features: Chat Interface, Streaming Tool Calls, Tool Detection, JSON Repair, Dual Response
- Document MCP Chat Interface Technical Details with streaming assembly specs
- Add intelligent tool detection system with keyword-based scheduling detection
- Document dual response architecture (Nebius API + MCP fallback)
- Update Future Work with MCP-specific enhancements and performance optimization
- Add streaming tool call processing details (200+ delta handling, 60s timeout)
- Update usage examples with specific technical tasks
- Document JSON repair system and comprehensive error handling mechanisms

README.md CHANGED
@@ -24,31 +24,26 @@ It takes a project description, breaks it down into actionable tasks through a [
24
  **Source Code on GitHub:**
25
  [https://github.com/blackopsrepl/yuga-planner](https://github.com/blackopsrepl/yuga-planner)
26
 
27
- ### Gradio Web Demo Usage
28
 
29
- 1. Go to [the live demo](https://huggingface.co/spaces/blackopsrepl/yuga-planner) or [http://localhost:7860](http://localhost:7860)
30
-
31
- 2. **Upload project files** or **use mock projects:**
32
- - Upload one or more Markdown project file(s), then click "Load Data"
33
- - OR select from pre-configured mock projects for quick testing
34
- - Each file will be taken as a separate project
35
- - The app will parse, decompose, and estimate tasks using LLM agents
36
-
37
- 3. **Generate schedule:**
38
- - Click "Solve" to generate an optimal schedule against a randomly generated team
39
- - View results interactively with real-time solver progress
40
- - Task order is preserved within each project
41
 
42
  ### MCP Tool Usage
43
-
44
  1. **In any MCP-compatible chatbot or agent platform:**
45
  ```
46
  use yuga-planner mcp tool
47
  Task Description: [Your task description]
48
  ```
49
-
50
  2. **Attach your calendar file (.ics)** to provide existing commitments
51
-
52
  3. **Receive optimized schedule** that integrates your new task with existing calendar events
53
 
54
  ## Architecture
@@ -61,6 +56,13 @@ Yuga Planner follows a **service-oriented architecture** with clear separation o
61
  - **StateService:** Centralized state management for job tracking and schedule storage
62
  - **LoggingService:** Real-time log streaming for UI feedback and debugging
63
  - **MockProjectService:** Provides sample project data for testing and demos
 
 
 
 
 
 
 
64
 
65
  ### System Components
66
  - **Gradio UI:** Modern web interface with real-time updates and interactive schedule visualization
@@ -86,17 +88,25 @@ Yuga Planner follows a **service-oriented architecture** with clear separation o
86
  | **Business Hours Enforcement** | Respects 9:00-18:00 working hours with lunch break exclusion | βœ… |
87
  | **Weekend Scheduling Prevention** | Hard constraint preventing weekend task assignments | βœ… |
88
  | **MCP Endpoint** | API endpoint for MCP tool integration with calendar support | βœ… |
 
 
 
 
 
89
 
90
  ## 🎯 Two Usage Modes
91
-
92
  Yuga Planner operates as **two separate systems** serving different use cases:
93
 
94
- ### 1. πŸ–₯️ Gradio Web Demo
95
- **Purpose:** Interactive team scheduling for project management
96
- - **Access:** [Live demo](https://huggingface.co/spaces/blackopsrepl/yuga-planner) or local web interface
97
- - **Input:** Upload Markdown project files or use pre-configured mock projects
98
- - **Team:** Schedules against a **randomly generated team** with diverse skills and availability
99
- - **Use Case:** Project managers scheduling real teams for complex multi-project workloads
 
 
 
 
100
 
101
  ### 2. πŸ€– MCP Personal Tool
102
  **Purpose:** Individual task scheduling integrated with personal calendars
@@ -115,7 +125,27 @@ Tool Response: Optimized schedule created - EC2 setup task assigned to
115
  available time slots around your existing meetings
116
  ```
117
 
118
- ## 🧩 MCP Tool Integration Details
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
119
 
120
  **Features:**
121
  - Accepts calendar files and user task descriptions via chat interface
@@ -125,105 +155,84 @@ available time slots around your existing meetings
125
  - Designed for seamless chatbot and agent workflow integration
126
 
127
  **Current Limitations:**
128
- - **Cross-system integration:** Gradio web demo and MCP personal tool operate as separate systems
129
  - **Multi-timezone support:** Currently operates in a single timezone context with UTC conversion for consistency. Calendar events from different timezones are normalized to the same scheduling context.
130
 
131
  See the [CHANGELOG.md](CHANGELOG.md) for details on recent MCP-related changes.
132
 
133
- ### Recent Improvements βœ…
134
-
135
- - **Service Architecture Refactoring:** Complete service-oriented architecture with proper encapsulation and clean boundaries
136
- - **State Management:** Centralized state handling through dedicated StateService
137
- - **Handler Compliance:** Clean separation between UI handlers and business logic services
138
- - **Method Encapsulation:** Fixed all private method violations for better code maintainability
139
-
140
  ### Work in Progress
141
-
142
  - **πŸ”§ Gradio UI overhaul:** Enhanced user experience and visual improvements
143
  - **πŸ” Migration to Pydantic models:** Type-safe data validation and serialization
144
- - **πŸ“š Migrate from violation_analyzer to Timefold dedicated libraries**
145
  - **⚑ Enhanced timezone support:** Multi-timezone calendar integration for international scheduling
146
 
147
  ### Future Work
148
-
149
  #### System Integration Roadmap
150
- Currently, the **Gradio web demo** and **MCP personal tool** operate as separate systems. As the project evolves, these will become **more integrated**, enabling:
151
  - **Unified scheduling engine** that can handle both team management and personal productivity in one interface
152
  - **Hybrid workflows** where personal tasks can be coordinated with team projects
153
  - **Cross-system data sharing** between web demo projects and personal MCP calendars
154
  - **Seamless switching** between team management and individual task planning modes
155
 
156
  #### Core Feature Enhancements
 
 
 
157
  - **RAG:** validation of task decomposition and estimation against industry relevant literature
158
- - **More granular task dependency:** representation of tasks in a tree instead of a list to allow overlap within projects, where feasible/convenient
159
  - **Input from GitHub issues:** instead of processing markdown directly, it creates a list by parsing issue
160
- - **Chat interface:** detection of user intent, with on-the-fly CRUD operations on team, tasks and schedules
161
  - **Reinforcement learning:** training the agent to improve task decomposition and estimation from GitHub history (e.g. diffs in timestamps, issue comments etc.)
162
 
163
  ## Prerequisites (Local/GitHub)
164
-
165
  - Python 3.10
166
  - Java 17+
167
  - Docker (optional, for containerized deployment)
168
  - Nebius API credentials (for LLM-powered features)
169
 
170
  ### Installation
171
-
172
  1. **Clone the repository:**
173
  ```bash
174
  git clone https://github.com/blackopsrepl/yuga-planner.git
175
  cd yuga-planner
176
  ```
177
-
178
  2. **Create a virtual environment:**
179
  ```bash
180
  make venv
181
  ```
182
-
183
  3. **Install dependencies:**
184
  ```bash
185
  make install
186
  ```
187
-
188
  4. **Set up environment variables / secrets:**
189
  ```bash
190
  make setup-secrets
191
  # Then edit tests/secrets/cred.py to add your API credentials
192
  ```
193
-
194
  5. **Run the app:**
195
  ```bash
196
  make run
197
  ```
198
 
199
  #### Docker (Local/GitHub)
200
-
201
  1. **Build the image:**
202
  ```bash
203
  docker build -t yuga-planner .
204
  ```
205
-
206
  2. **Run the container:**
207
  ```bash
208
- docker run -p 7860:786
209
  ```
210
 
211
  ---
212
 
213
  ## Python Dependencies
214
-
215
  See `requirements.txt` for full list.
216
 
217
  ---
218
 
219
  ## License
220
-
221
  This project is licensed under the Apache 2.0 License. See [LICENSE.txt](LICENSE.txt) for details.
222
 
223
  ---
224
 
225
  ## Acknowledgements
226
-
227
  - [Hugging Face](https://huggingface.co/)
228
  - [Gradio](https://gradio.app/)
229
  - [Nebius LLM](https://nebius.ai/)
 
24
  **Source Code on GitHub:**
25
  [https://github.com/blackopsrepl/yuga-planner](https://github.com/blackopsrepl/yuga-planner)
26
 
27
+ ### MCP Chat Interface Usage
28
 
29
+ 1. **Navigate to the Chat tab** in the Gradio interface
30
+ 2. **Upload your calendar file (.ics)** to provide existing commitments (optional)
31
+ 3. **Type your scheduling request** using natural language:
32
+ ```
33
+ "Create a new EC2 instance on AWS"
34
+ "Create a Svelte UI that allows me to query a postgresql database"
35
+ "Develop a chatbot UI based on Gradio"
36
+ ```
37
+ 4. **Receive intelligent responses** that combine conversational AI with task scheduling capabilities
38
+ 5. **View formatted schedules** with rich table output and status indicators
 
 
39
 
40
  ### MCP Tool Usage
 
41
  1. **In any MCP-compatible chatbot or agent platform:**
42
  ```
43
  use yuga-planner mcp tool
44
  Task Description: [Your task description]
45
  ```
 
46
  2. **Attach your calendar file (.ics)** to provide existing commitments
 
47
  3. **Receive optimized schedule** that integrates your new task with existing calendar events
48
 
49
  ## Architecture
 
56
  - **StateService:** Centralized state management for job tracking and schedule storage
57
  - **LoggingService:** Real-time log streaming for UI feedback and debugging
58
  - **MockProjectService:** Provides sample project data for testing and demos
59
+ - **MCPClientWrapper:** Manages MCP server functionality and tool definitions
60
+
61
+ ### MCP Integration Components
62
+ - **ToolCallAssembler:** Processes streaming tool call deltas from Nebius API into complete tool calls
63
+ - **ToolCallProcessor:** Executes completed tool calls via MCP backend with JSON repair functionality
64
+ - **Chat Interface:** Unified conversational AI + task scheduling with intelligent tool detection
65
+ - **Streaming Handler:** Real-time response processing with progress indicators and error recovery
66
 
67
  ### System Components
68
  - **Gradio UI:** Modern web interface with real-time updates and interactive schedule visualization
 
88
  | **Business Hours Enforcement** | Respects 9:00-18:00 working hours with lunch break exclusion | βœ… |
89
  | **Weekend Scheduling Prevention** | Hard constraint preventing weekend task assignments | βœ… |
90
  | **MCP Endpoint** | API endpoint for MCP tool integration with calendar support | βœ… |
91
+ | **Chat Interface with MCP** | Unified conversational AI + task scheduling interface | βœ… |
92
+ | **Streaming Tool Calls** | Real-time processing of tool calls from Nebius API | βœ… |
93
+ | **Intelligent Tool Detection** | Keyword-based detection for scheduling requests | βœ… |
94
+ | **JSON Repair & Recovery** | Robust handling of malformed streaming data | βœ… |
95
+ | **Dual Response System** | Nebius API with MCP fallback for reliability | βœ… |
96
 
97
  ## 🎯 Two Usage Modes
 
98
  Yuga Planner operates as **two separate systems** serving different use cases:
99
 
100
+ ### 1. πŸ’¬ MCP Chat Interface
101
+ **Purpose:** Conversational AI with integrated task scheduling capabilities
102
+ - **Access:** Chat tab in the Gradio interface
103
+ - **Input:** Natural language requests + optional `.ics` calendar files
104
+ - **Features:**
105
+ - Intelligent tool detection based on scheduling keywords
106
+ - Streaming responses with real-time tool call assembly
107
+ - Rich table formatting for schedule results
108
+ - Dual response system (Nebius API + MCP fallback)
109
+ - **Use Case:** Interactive scheduling through natural conversation
110
 
111
  ### 2. πŸ€– MCP Personal Tool
112
  **Purpose:** Individual task scheduling integrated with personal calendars
 
125
  available time slots around your existing meetings
126
  ```
127
 
128
+ ## 🧩 MCP Chat Interface Technical Details
129
+ ### Intelligent Tool Detection
130
+ The chat interface automatically detects scheduling requests using keyword analysis:
131
+ ```python
132
+ scheduling_keywords = [
133
+ 'schedule', 'task', 'calendar', 'plan', 'organize',
134
+ 'meeting', 'appointment', 'project', 'deadline',
135
+ 'create', 'setup', 'implement', 'develop'
136
+ ]
137
+ ```
138
+
139
+ ### Streaming Tool Call Processing
140
+ - **Delta Assembly:** Collects 200+ streaming deltas into complete tool calls
141
+ - **JSON Repair:** Handles malformed JSON from streaming responses
142
+ - **Progress Indicators:** Real-time feedback during tool processing
143
+ - **Timeout Protection:** 60-second timeout for MCP operations
144
+
145
+ ### Dual Response System
146
+ - **Primary:** Nebius API with tool calling capabilities
147
+ - **Fallback:** Direct MCP backend invocation when tool assembly fails
148
+ - **Error Recovery:** Comprehensive error handling and graceful degradation
149
 
150
  **Features:**
151
  - Accepts calendar files and user task descriptions via chat interface
 
155
  - Designed for seamless chatbot and agent workflow integration
156
 
157
  **Current Limitations:**
 
158
  - **Multi-timezone support:** Currently operates in a single timezone context with UTC conversion for consistency. Calendar events from different timezones are normalized to the same scheduling context.
159
 
160
  See the [CHANGELOG.md](CHANGELOG.md) for details on recent MCP-related changes.
161
 
 
 
 
 
 
 
 
162
  ### Work in Progress
 
163
  - **πŸ”§ Gradio UI overhaul:** Enhanced user experience and visual improvements
164
  - **πŸ” Migration to Pydantic models:** Type-safe data validation and serialization
 
165
  - **⚑ Enhanced timezone support:** Multi-timezone calendar integration for international scheduling
166
 
167
  ### Future Work
 
168
  #### System Integration Roadmap
 
169
  - **Unified scheduling engine** that can handle both team management and personal productivity in one interface
170
  - **Hybrid workflows** where personal tasks can be coordinated with team projects
171
  - **Cross-system data sharing** between web demo projects and personal MCP calendars
172
  - **Seamless switching** between team management and individual task planning modes
173
 
174
  #### Core Feature Enhancements
175
+ - **Multi-Tool Support:** Extend chat interface to support additional MCP tools beyond scheduling
176
+ - **Calendar Integration:** Direct calendar service integration (Google, Outlook)
177
+ - **Performance Optimization:** Enhanced streaming assembly for large tool calls
178
  - **RAG:** validation of task decomposition and estimation against industry relevant literature
179
+ - **More granular task dependency:** representation of tasks in a DAG instead of a list to allow overlap within projects, where feasible/convenient
180
  - **Input from GitHub issues:** instead of processing markdown directly, it creates a list by parsing issue
 
181
  - **Reinforcement learning:** training the agent to improve task decomposition and estimation from GitHub history (e.g. diffs in timestamps, issue comments etc.)
182
 
183
  ## Prerequisites (Local/GitHub)
 
184
  - Python 3.10
185
  - Java 17+
186
  - Docker (optional, for containerized deployment)
187
  - Nebius API credentials (for LLM-powered features)
188
 
189
  ### Installation
 
190
  1. **Clone the repository:**
191
  ```bash
192
  git clone https://github.com/blackopsrepl/yuga-planner.git
193
  cd yuga-planner
194
  ```
 
195
  2. **Create a virtual environment:**
196
  ```bash
197
  make venv
198
  ```
 
199
  3. **Install dependencies:**
200
  ```bash
201
  make install
202
  ```
 
203
  4. **Set up environment variables / secrets:**
204
  ```bash
205
  make setup-secrets
206
  # Then edit tests/secrets/cred.py to add your API credentials
207
  ```
 
208
  5. **Run the app:**
209
  ```bash
210
  make run
211
  ```
212
 
213
  #### Docker (Local/GitHub)
 
214
  1. **Build the image:**
215
  ```bash
216
  docker build -t yuga-planner .
217
  ```
 
218
  2. **Run the container:**
219
  ```bash
220
+ docker run -p 7860:7860 yuga-planner
221
  ```
222
 
223
  ---
224
 
225
  ## Python Dependencies
 
226
  See `requirements.txt` for full list.
227
 
228
  ---
229
 
230
  ## License
 
231
  This project is licensed under the Apache 2.0 License. See [LICENSE.txt](LICENSE.txt) for details.
232
 
233
  ---
234
 
235
  ## Acknowledgements
 
236
  - [Hugging Face](https://huggingface.co/)
237
  - [Gradio](https://gradio.app/)
238
  - [Nebius LLM](https://nebius.ai/)
src/app.py CHANGED
@@ -7,23 +7,10 @@ from utils.logging_config import setup_logging, get_logger
7
  setup_logging()
8
  logger = get_logger(__name__)
9
 
10
- from utils.load_secrets import load_secrets
11
-
12
- if not os.getenv("NEBIUS_API_KEY") or not os.getenv("NEBIUS_MODEL"):
13
- load_secrets("tests/secrets/creds.py")
14
-
15
-
16
- # from handlers.web_backend import (
17
- # load_data,
18
- # show_solved,
19
- # start_timer,
20
- # auto_poll,
21
- # show_mock_project_content,
22
- # )
23
-
24
  from handlers.mcp_backend import process_message_and_attached_file
 
 
25
 
26
- from services import MockProjectService
27
 
28
  # Store last chat message and file in global variables (for demo purposes)
29
  last_message_body = None
@@ -55,14 +42,11 @@ def app(debug: bool = False):
55
 
56
  **Yuga Planner** is a neuro-symbolic system that combines AI agents with constraint optimization
57
  for intelligent scheduling.
58
-
59
- ## πŸ”Œ **Using as MCP Tool**
60
-
61
- You can use Yuga Planner as an MCP server to integrate scheduling into your AI workflows.
62
  """
63
  )
64
 
65
- _draw_info_page(debug)
 
66
 
67
  # Register the MCP tool as an API endpoint
68
  gr.api(process_message_and_attached_file)
@@ -70,126 +54,6 @@ def app(debug: bool = False):
70
  return demo
71
 
72
 
73
- def _draw_info_page(debug: bool = False):
74
- with gr.Tab("πŸ“‹ Information"):
75
-
76
- def get_server_url():
77
- try:
78
- return gr.get_state().server_url + "/gradio_api/mcp/sse"
79
- except:
80
- return "https://blackopsrepl-yuga-planner.hf.space/gradio_api/mcp/sse"
81
-
82
- gr.Textbox(
83
- value=get_server_url(),
84
- label="🌐 MCP Server Endpoint",
85
- interactive=False,
86
- max_lines=1,
87
- )
88
-
89
- with gr.Accordion("πŸ“ MCP Setup Instructions", open=True):
90
- gr.Markdown(
91
- """
92
- ### 1. **Cursor Setup Instructions (should work from any MCP client!)**
93
-
94
- **For Cursor AI Editor:**
95
- 1. Create or edit your MCP configuration file: `~/.cursor/mcp.json`
96
- 2. Add the yuga-planner server configuration:
97
- ```json
98
- {
99
- "mcpServers": {
100
- "yuga-planner": {
101
- "url": -> "Insert the above endpoint URL here"
102
- }
103
- }
104
- }
105
- ```
106
- 3. If you already have other MCP servers, add `yuga-planner` to the existing `mcpServers` object
107
- 4. Restart Cursor to load the new configuration
108
- 5. The tool will be available in your chat
109
-
110
- ### 2. **Usage Example**
111
- """
112
- )
113
-
114
- gr.Textbox(
115
- value="""use yuga-planner mcp tool
116
- Task Description: Create a new EC2 instance on AWS
117
-
118
- [Attach your calendar.ics file to provide existing commitments]
119
-
120
- Tool Response: Optimized schedule created - EC2 setup task assigned to
121
- available time slots around your existing meetings
122
- [Returns JSON response with schedule data]
123
-
124
- User: show all fields as a table, ordered by start date
125
-
126
- [Displays formatted schedule table with all tasks and calendar events]""",
127
- label="πŸ’¬ Cursor Chat Usage Example",
128
- interactive=False,
129
- lines=10,
130
- )
131
-
132
- gr.Markdown(
133
- """
134
- ### 3. **What it does**
135
-
136
- **Personal Task Scheduling with Calendar Integration:**
137
-
138
- 1. πŸ“… **Parses your calendar** (.ics file) for existing commitments
139
- 2. πŸ€– **AI breaks down your task** into actionable subtasks using LLamaIndex + Nebius AI
140
- 3. ⚑ **Constraint-based optimization** finds optimal time slots around your existing schedule
141
- 4. πŸ“‹ **Returns complete solved schedule** integrated with your personal calendar events
142
- 5. πŸ•˜ **Respects business hours** (9:00-18:00) and excludes weekends automatically
143
- 6. πŸ“Š **JSON response format** - Ask to "show all fields as a table, ordered by start date" for readable formatting
144
-
145
- **Designed for**: Personal productivity and task planning around existing appointments in Cursor.
146
- """
147
- )
148
-
149
- if debug:
150
- with gr.Tab("πŸ› Debug Info"):
151
- gr.Markdown(
152
- """
153
- # πŸ› Debug Information
154
-
155
- **Debug Mode Enabled** - Additional system information and controls available.
156
- """
157
- )
158
-
159
- with gr.Accordion("πŸ”§ **Environment Details**", open=True):
160
- import os
161
-
162
- env_info = f"""
163
- **🐍 Python Environment**
164
- - Debug Mode: {debug}
165
- - YUGA_DEBUG: {os.getenv('YUGA_DEBUG', 'Not Set')}
166
- - Nebius API Key: {'βœ… Set' if os.getenv('NEBIUS_API_KEY') else '❌ Not Set'}
167
- - Nebius Model: {os.getenv('NEBIUS_MODEL', 'Not Set')}
168
-
169
- **🌐 Server Information**
170
- - MCP Endpoint: {get_server_url()}
171
- - Current Working Directory: {os.getcwd()}
172
- """
173
- gr.Markdown(env_info)
174
-
175
- with gr.Accordion("πŸ“Š **System Status**", open=False):
176
- gr.Markdown(
177
- """
178
- **πŸ”„ Service Status**
179
- - DataService: βœ… Active
180
- - ScheduleService: βœ… Active
181
- - StateService: βœ… Active
182
- - LoggingService: βœ… Active
183
- - MockProjectService: βœ… Active
184
-
185
- **πŸ”Œ Integration Status**
186
- - MCP Server: βœ… Enabled
187
- - Gradio API: βœ… Active
188
- - Real-time Logs: βœ… Streaming
189
- """
190
- )
191
-
192
-
193
  if __name__ == "__main__":
194
  parser = argparse.ArgumentParser(
195
  description="Yuga Planner - Team Scheduling Application"
 
7
  setup_logging()
8
  logger = get_logger(__name__)
9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
10
  from handlers.mcp_backend import process_message_and_attached_file
11
+ from ui.pages.chat import draw_chat_page
12
+ from ui.pages.mcp_info import draw_info_page
13
 
 
14
 
15
  # Store last chat message and file in global variables (for demo purposes)
16
  last_message_body = None
 
42
 
43
  **Yuga Planner** is a neuro-symbolic system that combines AI agents with constraint optimization
44
  for intelligent scheduling.
 
 
 
 
45
  """
46
  )
47
 
48
+ draw_chat_page(debug)
49
+ draw_info_page(debug)
50
 
51
  # Register the MCP tool as an API endpoint
52
  gr.api(process_message_and_attached_file)
 
54
  return demo
55
 
56
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
57
  if __name__ == "__main__":
58
  parser = argparse.ArgumentParser(
59
  description="Yuga Planner - Team Scheduling Application"
src/domain.py CHANGED
@@ -1,6 +1,12 @@
1
  import os, warnings
2
  from dataclasses import dataclass
3
 
 
 
 
 
 
 
4
  # =========================
5
  # MOCK PROJECTS
6
  # =========================
 
1
  import os, warnings
2
  from dataclasses import dataclass
3
 
4
+ # Load secrets if environment variables are not set
5
+ from utils.load_secrets import load_secrets
6
+
7
+ if not os.getenv("NEBIUS_API_KEY") or not os.getenv("NEBIUS_MODEL"):
8
+ load_secrets("tests/secrets/creds.py")
9
+
10
  # =========================
11
  # MOCK PROJECTS
12
  # =========================
src/handlers/__init__.py CHANGED
@@ -16,6 +16,12 @@ from .mcp_backend import (
16
  process_message_and_attached_file,
17
  )
18
 
 
 
 
 
 
 
19
  __all__ = [
20
  "load_data",
21
  "show_solved",
@@ -23,4 +29,7 @@ __all__ = [
23
  "auto_poll",
24
  "show_mock_project_content",
25
  "process_message_and_attached_file",
 
 
 
26
  ]
 
16
  process_message_and_attached_file,
17
  )
18
 
19
+ from .tool_call_handler import (
20
+ ToolCallAssembler,
21
+ ToolCallProcessor,
22
+ create_tool_call_handler,
23
+ )
24
+
25
  __all__ = [
26
  "load_data",
27
  "show_solved",
 
29
  "auto_poll",
30
  "show_mock_project_content",
31
  "process_message_and_attached_file",
32
+ "ToolCallAssembler",
33
+ "ToolCallProcessor",
34
+ "create_tool_call_handler",
35
  ]
src/handlers/tool_call_handler.py ADDED
@@ -0,0 +1,454 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import re
3
+ import asyncio
4
+ from typing import Dict, List, Any, Optional
5
+ from utils.logging_config import setup_logging, get_logger
6
+
7
+ # Initialize logging
8
+ setup_logging()
9
+ logger = get_logger(__name__)
10
+
11
+
12
+ class ToolCallAssembler:
13
+ """Handles streaming tool call assembly from API responses"""
14
+
15
+ def __init__(self):
16
+ self.tool_calls: Dict[int, Dict[str, Any]] = {}
17
+ self.reset()
18
+
19
+ def reset(self):
20
+ """Reset the assembler for a new conversation"""
21
+ self.tool_calls = {}
22
+
23
+ def process_delta(self, delta: Dict[str, Any]) -> None:
24
+ """Process a single delta from streaming response"""
25
+ if "tool_calls" not in delta:
26
+ return
27
+
28
+ for tool_call_delta in delta["tool_calls"]:
29
+ index = tool_call_delta.get("index", 0)
30
+
31
+ # Initialize tool call if not exists
32
+ if index not in self.tool_calls:
33
+ self.tool_calls[index] = {
34
+ "id": "",
35
+ "type": "function",
36
+ "function": {"name": "", "arguments": ""},
37
+ }
38
+
39
+ # Update tool call components
40
+ if "id" in tool_call_delta:
41
+ self.tool_calls[index]["id"] = tool_call_delta["id"]
42
+ if "type" in tool_call_delta:
43
+ self.tool_calls[index]["type"] = tool_call_delta["type"]
44
+ if "function" in tool_call_delta:
45
+ if "name" in tool_call_delta["function"]:
46
+ self.tool_calls[index]["function"]["name"] = tool_call_delta[
47
+ "function"
48
+ ]["name"]
49
+ if "arguments" in tool_call_delta["function"]:
50
+ # Append arguments (they come in chunks)
51
+ self.tool_calls[index]["function"]["arguments"] += tool_call_delta[
52
+ "function"
53
+ ]["arguments"]
54
+
55
+ def get_completed_tool_calls(self) -> List[Dict[str, Any]]:
56
+ """Get list of completed tool calls"""
57
+ completed = []
58
+ for tool_call in self.tool_calls.values():
59
+ # Check if tool call is complete (has name and valid JSON arguments)
60
+ if tool_call["function"]["name"] and tool_call["function"]["arguments"]:
61
+ try:
62
+ # Validate JSON arguments
63
+ json.loads(tool_call["function"]["arguments"])
64
+ completed.append(tool_call)
65
+ except json.JSONDecodeError as e:
66
+ logger.warning(
67
+ f"Tool call {tool_call['id']} has invalid JSON arguments: {e}"
68
+ )
69
+ logger.debug(f"Arguments: {tool_call['function']['arguments']}")
70
+
71
+ # Enhanced debugging for character 805 issue
72
+ args = tool_call["function"]["arguments"]
73
+ error_pos = getattr(e, "pos", 804) # Get error position
74
+
75
+ if error_pos > 0:
76
+ # Show context around the error
77
+ start = max(0, error_pos - 50)
78
+ end = min(len(args), error_pos + 50)
79
+ context = args[start:end]
80
+
81
+ logger.error(f"JSON Error Context (around char {error_pos}):")
82
+ logger.error(
83
+ f" Before error: '{args[max(0, error_pos-20):error_pos]}'"
84
+ )
85
+ logger.error(
86
+ f" At error: '{args[error_pos:error_pos+1] if error_pos < len(args) else 'END'}'"
87
+ )
88
+ logger.error(
89
+ f" After error: '{args[error_pos+1:error_pos+21] if error_pos < len(args) else ''}'"
90
+ )
91
+ logger.error(f" Full context: '{context}'")
92
+
93
+ # Check if it's the calendar data causing issues
94
+ if "calendar_file_content" in args:
95
+ # Find where calendar data starts and ends
96
+ cal_start = args.find('"calendar_file_content":"')
97
+ if cal_start != -1:
98
+ cal_data_start = cal_start + len(
99
+ '"calendar_file_content":"'
100
+ )
101
+ # Look for the closing quote
102
+ cal_end = args.find('"', cal_data_start + 1)
103
+ if cal_end != -1:
104
+ logger.error(
105
+ f"Calendar data length: {cal_end - cal_data_start}"
106
+ )
107
+ logger.error(
108
+ f"Calendar data starts at: {cal_data_start}"
109
+ )
110
+ logger.error(f"Calendar data ends at: {cal_end}")
111
+ logger.error(
112
+ f"Error position relative to cal data: {error_pos - cal_data_start}"
113
+ )
114
+ else:
115
+ logger.error("Calendar data has no closing quote!")
116
+ else:
117
+ logger.error(
118
+ "No calendar_file_content found in arguments"
119
+ )
120
+
121
+ # Try to repair the JSON by attempting common fixes
122
+ try:
123
+ repaired_args = self._attempt_json_repair(args)
124
+ if repaired_args:
125
+ json.loads(repaired_args) # Test if repair worked
126
+ logger.info(
127
+ f"Successfully repaired JSON for tool call {tool_call['id']}"
128
+ )
129
+ # Update the tool call with repaired arguments
130
+ tool_call["function"]["arguments"] = repaired_args
131
+ completed.append(tool_call)
132
+ continue
133
+ except:
134
+ logger.debug("JSON repair attempt failed")
135
+
136
+ return completed
137
+
138
+ def _attempt_json_repair(self, broken_json: str) -> str:
139
+ """Attempt to repair common JSON issues"""
140
+ try:
141
+ # Common issue: Missing closing brace
142
+ if not broken_json.strip().endswith("}"):
143
+ return broken_json.strip() + "}"
144
+
145
+ # Common issue: Malformed calendar data causing JSON breaks
146
+ # The error is likely around character 795 in the base64 data
147
+ if '"calendar_file_content":"' in broken_json:
148
+ # Try to find and fix the calendar field
149
+ start_pattern = '"calendar_file_content":"'
150
+ start_idx = broken_json.find(start_pattern)
151
+
152
+ if start_idx != -1:
153
+ content_start = start_idx + len(start_pattern)
154
+
155
+ # Find the actual end of the JSON by looking for the pattern
156
+ # that should follow calendar content: "} or ",
157
+ remaining = broken_json[content_start:]
158
+
159
+ # Look for what appears to be a duplicate task_description field
160
+ # This indicates where the JSON got corrupted
161
+ duplicate_pattern = r'"[\s\S]*?\{\s*"task_description"'
162
+ match = re.search(duplicate_pattern, remaining)
163
+
164
+ if match:
165
+ # The base64 data should end before this match
166
+ clean_end_pos = content_start + match.start()
167
+
168
+ # Extract the clean calendar data (everything before the corrupted part)
169
+ clean_calendar_data = (
170
+ broken_json[content_start:clean_end_pos]
171
+ .rstrip('"')
172
+ .rstrip()
173
+ )
174
+
175
+ # Extract the task description from the duplicate section
176
+ dup_start = content_start + match.start()
177
+ dup_content = broken_json[dup_start:]
178
+
179
+ # Find the actual JSON structure in the duplicate content
180
+ dup_json_start = dup_content.find('{"task_description"')
181
+ if dup_json_start != -1:
182
+ clean_rest = dup_content[dup_json_start:]
183
+
184
+ # Remove any trailing garbage that might cause JSON issues
185
+ if not clean_rest.endswith("}"):
186
+ # Try to find the proper ending
187
+ brace_count = 0
188
+ proper_end = -1
189
+ for i, char in enumerate(clean_rest):
190
+ if char == "{":
191
+ brace_count += 1
192
+ elif char == "}":
193
+ brace_count -= 1
194
+ if brace_count == 0:
195
+ proper_end = i + 1
196
+ break
197
+
198
+ if proper_end != -1:
199
+ clean_rest = clean_rest[:proper_end]
200
+ else:
201
+ clean_rest = clean_rest.rstrip() + "}"
202
+
203
+ # Parse the clean rest to extract just the fields we need
204
+ try:
205
+ rest_obj = json.loads(clean_rest)
206
+ task_desc = rest_obj.get("task_description", "")
207
+
208
+ # Reconstruct the proper JSON
209
+ repaired_obj = {
210
+ "task_description": task_desc,
211
+ "calendar_file_content": clean_calendar_data,
212
+ }
213
+
214
+ repaired = json.dumps(repaired_obj)
215
+ logger.debug(
216
+ f"Successfully reconstructed JSON with task: {task_desc[:50]}..."
217
+ )
218
+ logger.debug(
219
+ f"Calendar data length: {len(clean_calendar_data)}"
220
+ )
221
+ return repaired
222
+
223
+ except json.JSONDecodeError as e:
224
+ logger.debug(
225
+ f"Failed to parse extracted duplicate section: {e}"
226
+ )
227
+ # Fallback to simple reconstruction
228
+ repaired = f'{{"task_description":"","calendar_file_content":"{clean_calendar_data}"}}'
229
+ logger.debug(
230
+ f"Fallback repair with calendar data length: {len(clean_calendar_data)}"
231
+ )
232
+ return repaired
233
+
234
+ # Alternative: Look for more traditional closing patterns
235
+ # Find where the base64 data properly ends with quote+comma or quote+brace
236
+ for i, char in enumerate(remaining):
237
+ if char == '"':
238
+ # Check what follows the quote
239
+ next_chars = (
240
+ remaining[i : i + 3]
241
+ if i + 3 <= len(remaining)
242
+ else remaining[i:]
243
+ )
244
+ if next_chars.startswith('",') or next_chars.startswith(
245
+ '"}'
246
+ ):
247
+ # This looks like a proper end
248
+ calendar_data = remaining[:i]
249
+ rest = remaining[i:]
250
+ repaired = (
251
+ broken_json[:content_start] + calendar_data + rest
252
+ )
253
+ logger.debug(
254
+ f"Pattern-based repair: calendar data length {len(calendar_data)}"
255
+ )
256
+ return repaired
257
+
258
+ # If calendar repair didn't work, try other common fixes
259
+ # Remove any stray characters that might have been inserted
260
+ cleaned = re.sub(
261
+ r"[^\x20-\x7E]", "", broken_json
262
+ ) # Remove non-printable chars
263
+ if cleaned != broken_json:
264
+ logger.debug("Attempted to remove non-printable characters")
265
+ return cleaned
266
+
267
+ # Last resort: Try to extract valid JSON from the string
268
+ # Look for the first { and try to find matching }
269
+ first_brace = broken_json.find("{")
270
+ if first_brace != -1:
271
+ brace_count = 0
272
+ for i in range(first_brace, len(broken_json)):
273
+ if broken_json[i] == "{":
274
+ brace_count += 1
275
+ elif broken_json[i] == "}":
276
+ brace_count -= 1
277
+ if brace_count == 0:
278
+ candidate = broken_json[first_brace : i + 1]
279
+ try:
280
+ json.loads(candidate)
281
+ logger.debug(
282
+ f"Extracted valid JSON segment of length {len(candidate)}"
283
+ )
284
+ return candidate
285
+ except:
286
+ continue
287
+
288
+ return None
289
+ except Exception as e:
290
+ logger.debug(f"JSON repair attempt failed: {e}")
291
+ return None
292
+
293
+ def debug_info(self) -> Dict[str, Any]:
294
+ """Get debug information about current tool calls"""
295
+ info = {
296
+ "total_tool_calls": len(self.tool_calls),
297
+ "completed_tool_calls": len(self.get_completed_tool_calls()),
298
+ "tool_calls_detail": {},
299
+ }
300
+
301
+ for index, tool_call in self.tool_calls.items():
302
+ info["tool_calls_detail"][index] = {
303
+ "id": tool_call["id"],
304
+ "function_name": tool_call["function"]["name"],
305
+ "arguments_length": len(tool_call["function"]["arguments"]),
306
+ "arguments_preview": tool_call["function"]["arguments"][:100] + "..."
307
+ if len(tool_call["function"]["arguments"]) > 100
308
+ else tool_call["function"]["arguments"],
309
+ "is_json_valid": self._is_valid_json(
310
+ tool_call["function"]["arguments"]
311
+ ),
312
+ }
313
+
314
+ return info
315
+
316
+ def _is_valid_json(self, json_str: str) -> bool:
317
+ """Check if string is valid JSON"""
318
+ try:
319
+ json.loads(json_str)
320
+ return True
321
+ except:
322
+ return False
323
+
324
+
325
+ class ToolCallProcessor:
326
+ """Processes completed tool calls"""
327
+
328
+ def __init__(self, mcp_client):
329
+ self.mcp_client = mcp_client
330
+ self.loop = None
331
+ try:
332
+ self.loop = asyncio.get_event_loop()
333
+ except RuntimeError:
334
+ self.loop = asyncio.new_event_loop()
335
+ asyncio.set_event_loop(self.loop)
336
+
337
+ def process_tool_calls(self, tool_calls: List[Dict[str, Any]], message: str) -> str:
338
+ """Process a list of tool calls and return response text"""
339
+ if not tool_calls:
340
+ return ""
341
+
342
+ response_parts = []
343
+
344
+ for tool_call in tool_calls:
345
+ logger.info(f"Processing tool call: {tool_call}")
346
+
347
+ try:
348
+ # Extract function details
349
+ function_name = tool_call.get("function", {}).get("name", "")
350
+ function_args_str = tool_call.get("function", {}).get("arguments", "{}")
351
+
352
+ if function_name == "schedule_tasks_with_calendar":
353
+ result = self._process_scheduling_tool(function_args_str, message)
354
+ response_parts.append(result)
355
+ else:
356
+ logger.debug(f"Ignoring non-scheduling tool: {function_name}")
357
+
358
+ except Exception as e:
359
+ logger.error(f"Error processing tool call: {e}")
360
+ response_parts.append(f"\n\n❌ **Error processing tool call:** {str(e)}")
361
+
362
+ return "".join(response_parts)
363
+
364
+ def _process_scheduling_tool(self, function_args_str: str, message: str) -> str:
365
+ """Process scheduling tool call"""
366
+ try:
367
+ # Parse arguments
368
+ args = json.loads(function_args_str)
369
+ task_description = args.get("task_description", "")
370
+ calendar_content = args.get("calendar_file_content", "none")
371
+
372
+ # Extract calendar data from message if available (override args)
373
+ calendar_match = re.search(r"\[CALENDAR_DATA:([^\]]+)\]", message)
374
+ if calendar_match:
375
+ calendar_content = calendar_match.group(1)
376
+ logger.debug("Found calendar data in message, overriding args")
377
+
378
+ logger.info(f"Calling MCP scheduling tool with task: {task_description}")
379
+
380
+ # Call the scheduling tool
381
+ result = self.loop.run_until_complete(
382
+ self.mcp_client.call_scheduling_tool(task_description, calendar_content)
383
+ )
384
+
385
+ logger.info(
386
+ f"MCP tool completed with status: {result.get('status', 'unknown')}"
387
+ )
388
+
389
+ return self._format_scheduling_result(result, task_description)
390
+
391
+ except json.JSONDecodeError as e:
392
+ logger.error(f"JSON decode error in tool arguments: {e}")
393
+ logger.debug(f"Raw arguments: {function_args_str}")
394
+ return f"\n\n❌ **Error parsing tool arguments:** {str(e)}\n\nRaw arguments preview: {function_args_str[:200]}..."
395
+ except Exception as e:
396
+ logger.error(f"Tool execution error: {e}")
397
+ return f"\n\n❌ **Error processing scheduling request:** {str(e)}"
398
+
399
+ def _format_scheduling_result(
400
+ self, result: Dict[str, Any], task_description: str
401
+ ) -> str:
402
+ """Format the scheduling result for display"""
403
+ if result.get("status") == "success":
404
+ schedule = result.get("schedule", [])
405
+ calendar_entries = result.get("calendar_entries", [])
406
+
407
+ return f"""
408
+
409
+ πŸ“… **Schedule Generated Successfully!**
410
+
411
+ **Task:** {task_description}
412
+ **Calendar Events Processed:** {len(calendar_entries)}
413
+ **Total Scheduled Items:** {len(schedule)}
414
+
415
+ **Summary:**
416
+ - βœ… Existing calendar events preserved at original times
417
+ - πŸ†• New tasks optimized around existing commitments
418
+ - ⏰ All scheduling respects business hours (9:00-18:00)
419
+ - πŸ“‹ Complete schedule integration
420
+
421
+ To see the detailed schedule, ask me to "show the schedule as a table" or "format the schedule results".
422
+ """
423
+ elif result.get("status") == "timeout":
424
+ return f"""
425
+
426
+ ⏰ **Scheduling Analysis In Progress**
427
+
428
+ The schedule optimizer is still working on your complex task: "{task_description}"
429
+
430
+ This indicates a sophisticated scheduling challenge with multiple constraints. The system is finding the optimal arrangement for your tasks around existing calendar commitments.
431
+
432
+ You can check back in a few moments or try with a simpler task description.
433
+ """
434
+ else:
435
+ error_msg = result.get("error", "Unknown error")
436
+ return f"""
437
+
438
+ ❌ **Scheduling Error**
439
+
440
+ I encountered an issue while processing your scheduling request: {error_msg}
441
+
442
+ Please try:
443
+ - Simplifying your task description
444
+ - Checking if you have calendar conflicts
445
+ - Ensuring your .ics file is valid (if uploaded)
446
+ """
447
+
448
+
449
+ def create_tool_call_handler(mcp_client):
450
+ """Factory function to create a complete tool call handler"""
451
+ assembler = ToolCallAssembler()
452
+ processor = ToolCallProcessor(mcp_client)
453
+
454
+ return assembler, processor
src/services/__init__.py CHANGED
@@ -9,6 +9,7 @@ from .schedule import ScheduleService
9
  from .data import DataService
10
  from .mock_projects import MockProjectService
11
  from .state import StateService
 
12
 
13
  __all__ = [
14
  "LoggingService",
@@ -16,4 +17,5 @@ __all__ = [
16
  "DataService",
17
  "MockProjectService",
18
  "StateService",
 
19
  ]
 
9
  from .data import DataService
10
  from .mock_projects import MockProjectService
11
  from .state import StateService
12
+ from .mcp_client import MCPClientService
13
 
14
  __all__ = [
15
  "LoggingService",
 
17
  "DataService",
18
  "MockProjectService",
19
  "StateService",
20
+ "MCPClientService",
21
  ]
src/services/mcp_client.py ADDED
@@ -0,0 +1,80 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ MCP Client Service for Yuga Planner
3
+
4
+ This service handles Model Context Protocol client operations for scheduling tools.
5
+ Provides a clean interface for integrating MCP scheduling functionality.
6
+ """
7
+
8
+ import base64
9
+ from typing import Dict, Any
10
+
11
+ from utils.logging_config import setup_logging, get_logger
12
+ from handlers.mcp_backend import process_message_and_attached_file
13
+
14
+ # Initialize logging
15
+ setup_logging()
16
+ logger = get_logger(__name__)
17
+
18
+
19
+ class MCPClientService:
20
+ """Service for MCP client operations and scheduling tool integration"""
21
+
22
+ def __init__(self):
23
+ self.tools = [
24
+ {
25
+ "type": "function",
26
+ "function": {
27
+ "name": "schedule_tasks_with_calendar",
28
+ "description": "Create an optimized schedule by analyzing calendar events and breaking down tasks. Upload a calendar .ics file and describe the task you want to schedule.",
29
+ "parameters": {
30
+ "type": "object",
31
+ "properties": {
32
+ "task_description": {
33
+ "type": "string",
34
+ "description": "Description of the task or project to schedule (e.g., 'Create a new EC2 instance on AWS')",
35
+ },
36
+ "calendar_file_content": {
37
+ "type": "string",
38
+ "description": "Base64 encoded content of the .ics calendar file, or 'none' if no calendar provided",
39
+ },
40
+ },
41
+ "required": ["task_description", "calendar_file_content"],
42
+ },
43
+ },
44
+ }
45
+ ]
46
+
47
+ async def call_scheduling_tool(
48
+ self, task_description: str, calendar_file_content: str
49
+ ) -> Dict[str, Any]:
50
+ """
51
+ Call the scheduling backend tool.
52
+
53
+ Args:
54
+ task_description: Description of the task to schedule
55
+ calendar_file_content: Base64 encoded calendar content or 'none'
56
+
57
+ Returns:
58
+ Dict containing the scheduling result
59
+ """
60
+ try:
61
+ if calendar_file_content.lower() == "none":
62
+ file_content = b""
63
+ else:
64
+ file_content = base64.b64decode(calendar_file_content)
65
+
66
+ logger.debug(f"Calling MCP backend with task: {task_description}")
67
+ result = await process_message_and_attached_file(
68
+ file_content=file_content,
69
+ message_body=task_description,
70
+ file_name="calendar.ics",
71
+ )
72
+
73
+ logger.debug(
74
+ f"MCP backend returned status: {result.get('status', 'unknown')}"
75
+ )
76
+ return result
77
+
78
+ except Exception as e:
79
+ logger.error(f"Error calling scheduling tool: {e}")
80
+ return {"error": str(e), "status": "failed"}
src/ui/pages/chat.py ADDED
@@ -0,0 +1,705 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os, json, re
2
+ import gradio as gr
3
+ import asyncio
4
+ from typing import List, Dict, Any, Union, Generator
5
+ from contextlib import AsyncExitStack
6
+ from datetime import datetime, date
7
+
8
+ import requests
9
+
10
+ from handlers.mcp_backend import process_message_and_attached_file
11
+ from handlers.tool_call_handler import create_tool_call_handler
12
+ from services.mcp_client import MCPClientService
13
+
14
+ from utils.load_secrets import load_secrets
15
+
16
+ if not os.getenv("NEBIUS_API_KEY") or not os.getenv("NEBIUS_MODEL"):
17
+ load_secrets("tests/secrets/creds.py")
18
+
19
+ nebius_api_key = os.getenv("NEBIUS_API_KEY")
20
+ nebius_model = os.getenv("NEBIUS_MODEL")
21
+
22
+ from utils.logging_config import (
23
+ setup_logging,
24
+ get_logger,
25
+ start_session_logging,
26
+ get_session_logs,
27
+ )
28
+
29
+ # Initialize logging
30
+ setup_logging()
31
+ logger = get_logger(__name__)
32
+
33
+ # Global MCP client for the chat page
34
+ _mcp_client = None
35
+ _tool_assembler = None
36
+ _tool_processor = None
37
+
38
+ # Get or create event loop for MCP operations
39
+ try:
40
+ loop = asyncio.get_event_loop()
41
+ except RuntimeError:
42
+ loop = asyncio.new_event_loop()
43
+ asyncio.set_event_loop(loop)
44
+
45
+
46
+ class DateTimeEncoder(json.JSONEncoder):
47
+ """Custom JSON encoder that handles datetime objects and other non-serializable types"""
48
+
49
+ def default(self, obj):
50
+ if isinstance(obj, (datetime, date)):
51
+ return obj.isoformat()
52
+ elif hasattr(obj, "__dict__"):
53
+ return str(obj)
54
+ elif hasattr(obj, "to_dict"):
55
+ return obj.to_dict()
56
+ else:
57
+ return str(obj)
58
+
59
+
60
+ def safe_json_dumps(obj, **kwargs):
61
+ """Safely serialize objects to JSON, handling datetime and other non-serializable types"""
62
+ try:
63
+ return json.dumps(obj, cls=DateTimeEncoder, **kwargs)
64
+ except Exception as e:
65
+ logger.warning(
66
+ f"JSON serialization failed: {e}, falling back to string representation"
67
+ )
68
+ return json.dumps(
69
+ {"error": f"Serialization failed: {str(e)}", "raw_data": str(obj)[:1000]},
70
+ **kwargs,
71
+ )
72
+
73
+
74
+ def draw_chat_page(debug: bool = False):
75
+ logger.info(f"NEBIUS_MODEL: {nebius_model}")
76
+ logger.info(f"NEBIUS_API_KEY: {'Set' if nebius_api_key else 'Not Set'}")
77
+
78
+ if not nebius_model or not nebius_api_key:
79
+ logger.error(
80
+ "NEBIUS_MODEL or NEBIUS_API_KEY not found in environment variables"
81
+ )
82
+
83
+ with gr.Tab("πŸ’¬ Chat"):
84
+ gr.Markdown(
85
+ """
86
+ # πŸ’¬ Chat with Yuga Planner
87
+
88
+ This chatbot can help you with general questions and also schedule tasks around your calendar.
89
+ To use scheduling features, just describe your task and optionally upload a calendar .ics file.
90
+ """
91
+ )
92
+
93
+ if not nebius_model or not nebius_api_key:
94
+ gr.Markdown(
95
+ """
96
+ ⚠️ **Chat unavailable**: NEBIUS_MODEL or NEBIUS_API_KEY environment variables are not set.
97
+
98
+ Please configure your Nebius credentials to use the chat feature.
99
+ """
100
+ )
101
+ return
102
+
103
+ # Initialize MCP client and tool handlers as globals for this page
104
+ global _mcp_client, _tool_assembler, _tool_processor
105
+ _mcp_client = MCPClientService()
106
+ _tool_assembler, _tool_processor = create_tool_call_handler(_mcp_client)
107
+
108
+ # Create chat interface components
109
+ chatbot, msg, clear, calendar_file = create_chat_interface()
110
+
111
+ # Create parameter controls
112
+ (
113
+ system_message,
114
+ max_tokens_slider,
115
+ temperature_slider,
116
+ top_p_slider,
117
+ ) = create_chatbot_parameters()
118
+
119
+ msg.submit(
120
+ user_message, [msg, chatbot, calendar_file], [msg, chatbot], queue=False
121
+ ).then(
122
+ bot_response,
123
+ [
124
+ chatbot,
125
+ system_message,
126
+ max_tokens_slider,
127
+ temperature_slider,
128
+ top_p_slider,
129
+ ],
130
+ chatbot,
131
+ show_progress=True,
132
+ )
133
+
134
+ clear.click(lambda: [], None, chatbot, queue=False)
135
+
136
+
137
+ def create_chat_interface() -> tuple[gr.Chatbot, gr.Textbox, gr.Button, gr.File]:
138
+ """Create and return the chat interface components"""
139
+ chatbot = gr.Chatbot(type="messages")
140
+ msg = gr.Textbox(
141
+ label="Your message",
142
+ placeholder="Type your message here... For scheduling, describe your task and optionally upload a calendar file.",
143
+ )
144
+ calendar_file = gr.File(
145
+ label="πŸ“… Calendar File (.ics)", file_types=[".ics"], visible=True
146
+ )
147
+ clear = gr.Button("Clear")
148
+ return chatbot, msg, clear, calendar_file
149
+
150
+
151
+ def create_chatbot_parameters() -> tuple[gr.Textbox, gr.Slider, gr.Slider, gr.Slider]:
152
+ """Create and return the chatbot parameter controls"""
153
+ with gr.Accordion("Chatbot Parameters", open=False):
154
+ system_message = gr.Textbox(
155
+ value="You are a friendly and helpful AI assistant that specializes in task scheduling and productivity. You can help users plan and organize their work around existing calendar commitments. When users ask about scheduling tasks or mention calendar-related activities, you should use the schedule_tasks_with_calendar tool to create optimized schedules. If you see [Calendar file uploaded:] in a message, the user has provided calendar data that should be used for scheduling. Always use the scheduling tool when users mention tasks, projects, scheduling, planning, or similar requests.",
156
+ label="System message",
157
+ )
158
+ max_tokens_slider = gr.Slider(
159
+ minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"
160
+ )
161
+ temperature_slider = gr.Slider(
162
+ minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"
163
+ )
164
+ top_p_slider = gr.Slider(
165
+ minimum=0.1,
166
+ maximum=1.0,
167
+ value=0.95,
168
+ step=0.05,
169
+ label="Top-p (nucleus sampling)",
170
+ )
171
+ return system_message, max_tokens_slider, temperature_slider, top_p_slider
172
+
173
+
174
+ def user_message(message, history, calendar_file_obj):
175
+ # Handle calendar file upload
176
+ enhanced_message = message
177
+ if calendar_file_obj is not None:
178
+ # Read and encode the calendar file
179
+ try:
180
+ import base64
181
+
182
+ with open(calendar_file_obj, "rb") as f:
183
+ file_content = f.read()
184
+
185
+ encoded_content = base64.b64encode(file_content).decode("utf-8")
186
+ enhanced_message += (
187
+ f"\n\n[Calendar file uploaded: {calendar_file_obj.name}]"
188
+ )
189
+ enhanced_message += f"\n[CALENDAR_DATA:{encoded_content}]"
190
+ except Exception as e:
191
+ logger.error(f"Error reading calendar file: {e}")
192
+ enhanced_message += f"\n\n[Calendar file upload failed: {str(e)}]"
193
+
194
+ return "", history + [{"role": "user", "content": enhanced_message}]
195
+
196
+
197
+ def bot_response(history, system_message, max_tokens, temperature, top_p):
198
+ if not history:
199
+ return history
200
+
201
+ # Convert messages format to tuples for the respond function
202
+ history_tuples = []
203
+ for msg in history[:-1]: # All but the last message
204
+ if msg["role"] == "user":
205
+ history_tuples.append([msg["content"], ""])
206
+ elif msg["role"] == "assistant":
207
+ if history_tuples:
208
+ history_tuples[-1][1] = msg["content"]
209
+ else:
210
+ history_tuples.append(["", msg["content"]])
211
+
212
+ # Get the last user message
213
+ user_msg = history[-1]["content"]
214
+
215
+ logger.info(f"Bot response called with user message: {user_msg[:100]}...")
216
+
217
+ try:
218
+ # Get the response generator
219
+ response_gen = respond(
220
+ user_msg,
221
+ history_tuples,
222
+ system_message,
223
+ max_tokens,
224
+ temperature,
225
+ top_p,
226
+ )
227
+
228
+ # Stream responses to show progress - this is a generator function now
229
+ for response_chunk in response_gen:
230
+ updated_history = history.copy()
231
+ updated_history[-1] = {"role": "assistant", "content": response_chunk}
232
+ yield updated_history
233
+
234
+ except Exception as e:
235
+ logger.error(f"Error in bot_response: {e}")
236
+ import traceback
237
+
238
+ logger.error(f"Full traceback: {traceback.format_exc()}")
239
+ error_history = history.copy()
240
+ error_history[-1] = {"role": "assistant", "content": f"Error: {str(e)}"}
241
+ yield error_history
242
+
243
+
244
+ def respond(
245
+ message,
246
+ history: list[tuple[str, str]],
247
+ system_message,
248
+ max_tokens,
249
+ temperature,
250
+ top_p,
251
+ ) -> Generator[str, None, None]:
252
+ try:
253
+ # Start capturing logs for this session
254
+ start_session_logging()
255
+
256
+ # Reset tool assembler for new conversation
257
+ _tool_assembler.reset()
258
+
259
+ messages = [{"role": "system", "content": system_message}]
260
+
261
+ for val in history:
262
+ if val[0]:
263
+ messages.append({"role": "user", "content": val[0]})
264
+
265
+ if val[1]:
266
+ messages.append({"role": "assistant", "content": val[1]})
267
+
268
+ messages.append({"role": "user", "content": message})
269
+
270
+ # Check if this looks like a scheduling request
271
+ scheduling_keywords = [
272
+ "schedule",
273
+ "task",
274
+ "calendar",
275
+ "plan",
276
+ "organize",
277
+ "meeting",
278
+ "appointment",
279
+ "project",
280
+ "deadline",
281
+ "create",
282
+ "setup",
283
+ "implement",
284
+ "develop",
285
+ ]
286
+
287
+ is_scheduling_request = any(
288
+ keyword in message.lower() for keyword in scheduling_keywords
289
+ )
290
+
291
+ logger.info(f"Message: {message}")
292
+ logger.info(f"Is scheduling request: {is_scheduling_request}")
293
+
294
+ # Prepare payload for Nebius API
295
+ payload = {
296
+ "model": nebius_model,
297
+ "messages": messages,
298
+ "max_tokens": max_tokens,
299
+ "temperature": temperature,
300
+ "top_p": top_p,
301
+ "stream": True,
302
+ }
303
+
304
+ # Add tools if this might be a scheduling request
305
+ if is_scheduling_request:
306
+ logger.info("Adding tools to payload")
307
+ payload["tools"] = _mcp_client.tools
308
+ payload["tool_choice"] = "auto"
309
+ logger.debug(f"Tools payload: {_mcp_client.tools}")
310
+
311
+ else:
312
+ logger.info("No scheduling detected, not adding tools")
313
+
314
+ headers = {
315
+ "Authorization": f"Bearer {nebius_api_key}",
316
+ "Content-Type": "application/json",
317
+ }
318
+
319
+ logger.info(
320
+ f"Sending request to Nebius API with tools: {is_scheduling_request}"
321
+ )
322
+ logger.debug(f"Full payload: {json.dumps(payload, indent=2)}")
323
+
324
+ response = requests.post(
325
+ "https://api.studio.nebius.ai/v1/chat/completions",
326
+ headers=headers,
327
+ json=payload,
328
+ stream=True,
329
+ )
330
+
331
+ if response.status_code != 200:
332
+ logger.error(f"API error: {response.status_code} - {response.text}")
333
+ yield f"Error: API returned {response.status_code}: {response.text}"
334
+ return
335
+
336
+ response_text = ""
337
+
338
+ # Initial yield to show streaming is working
339
+ if is_scheduling_request:
340
+ yield "πŸ€– **Processing your scheduling request...**"
341
+
342
+ for line in response.iter_lines():
343
+ if line:
344
+ line = line.decode("utf-8")
345
+ if line.startswith("data: "):
346
+ data = line[6:] # Remove 'data: ' prefix
347
+ if data.strip() == "[DONE]":
348
+ break
349
+ try:
350
+ chunk = json.loads(data)
351
+ logger.debug(f"Received chunk: {chunk}")
352
+ if "choices" in chunk and len(chunk["choices"]) > 0:
353
+ delta = chunk["choices"][0].get("delta", {})
354
+ content = delta.get("content", "")
355
+ if content:
356
+ response_text += content
357
+ # For scheduling requests, include essential logs inline
358
+ if is_scheduling_request:
359
+ session_logs = get_session_logs()
360
+ if session_logs:
361
+ # Show only new logs since last yield
362
+ latest_logs = (
363
+ session_logs[-3:]
364
+ if len(session_logs) > 3
365
+ else session_logs
366
+ )
367
+ logs_text = "\n".join(
368
+ f" {log}" for log in latest_logs
369
+ )
370
+ yield response_text + f"\n\n{logs_text}"
371
+ else:
372
+ yield response_text
373
+ else:
374
+ yield response_text
375
+
376
+ # Process tool calls using our new handler
377
+ _tool_assembler.process_delta(delta)
378
+
379
+ except json.JSONDecodeError as e:
380
+ logger.error(f"JSON decode error: {e} for line: {line}")
381
+ continue
382
+
383
+ # Get completed tool calls and process them
384
+ completed_tool_calls = _tool_assembler.get_completed_tool_calls()
385
+
386
+ # Log debug info
387
+ debug_info = _tool_assembler.debug_info()
388
+ logger.info(f"Tool call assembly completed: {debug_info}")
389
+
390
+ if completed_tool_calls:
391
+ logger.info(f"Processing {len(completed_tool_calls)} completed tool calls")
392
+ yield response_text + "\n\nπŸ”§ **Processing scheduling request...**"
393
+
394
+ # Process tool calls using our new processor
395
+ tool_response = _tool_processor.process_tool_calls(
396
+ completed_tool_calls, message
397
+ )
398
+ response_text += tool_response
399
+ yield response_text
400
+
401
+ else:
402
+ logger.info("No completed tool calls found")
403
+ if is_scheduling_request:
404
+ logger.warning(
405
+ "Scheduling request detected but no completed tool calls"
406
+ )
407
+
408
+ # Log detailed debug info for troubleshooting
409
+ logger.error(f"Tool assembly debug info: {debug_info}")
410
+
411
+ yield response_text + "\n\n⚠️ **Scheduling request detected but tool not triggered or incomplete. Let me try calling the scheduler directly...**"
412
+
413
+ # Directly call the scheduling tool as fallback
414
+ try:
415
+ # Extract task description from message
416
+ task_description = message
417
+ calendar_content = "none"
418
+
419
+ # Extract calendar data if available
420
+ calendar_match = re.search(r"\[CALENDAR_DATA:([^\]]+)\]", message)
421
+ if calendar_match:
422
+ calendar_content = calendar_match.group(1)
423
+
424
+ # Show essential task processing logs inline
425
+ session_logs = get_session_logs()
426
+ processing_status = ""
427
+ if session_logs:
428
+ latest_logs = (
429
+ session_logs[-2:] if len(session_logs) > 2 else session_logs
430
+ )
431
+ processing_status = "\n" + "\n".join(
432
+ f" {log}" for log in latest_logs
433
+ )
434
+
435
+ yield response_text + f"\n\nπŸ”§ **Direct scheduling call for: {task_description}**\n⏳ *Processing...*{processing_status}"
436
+
437
+ logger.info("About to call MCP scheduling tool directly")
438
+
439
+ # Add timeout to prevent hanging
440
+ import asyncio
441
+ import concurrent.futures
442
+
443
+ def call_with_timeout():
444
+ try:
445
+ return loop.run_until_complete(
446
+ asyncio.wait_for(
447
+ _mcp_client.call_scheduling_tool(
448
+ task_description, calendar_content
449
+ ),
450
+ timeout=60.0, # 60 second timeout
451
+ )
452
+ )
453
+ except asyncio.TimeoutError:
454
+ return {
455
+ "error": "Timeout after 60 seconds",
456
+ "status": "timeout",
457
+ }
458
+
459
+ # Show progress during processing with essential logs
460
+ session_logs = get_session_logs()
461
+ analysis_status = ""
462
+ if session_logs:
463
+ latest_logs = (
464
+ session_logs[-3:] if len(session_logs) > 3 else session_logs
465
+ )
466
+ analysis_status = "\n" + "\n".join(
467
+ f" {log}" for log in latest_logs
468
+ )
469
+
470
+ yield response_text + f"\n\nπŸ”§ **Direct scheduling call for: {task_description}**\n⏳ *Analyzing calendar and generating tasks...*{analysis_status}"
471
+
472
+ try:
473
+ result = call_with_timeout()
474
+ except Exception as timeout_err:
475
+ logger.error(
476
+ f"MCP scheduling tool timed out or failed: {timeout_err}"
477
+ )
478
+ tool_response = f"\n\n⏰ **Scheduling timed out** - The request took longer than expected. Please try with a simpler task description."
479
+ response_text += tool_response
480
+ logger.info("Added timeout message to response")
481
+ yield response_text
482
+ else:
483
+ # Show progress for result processing
484
+ yield response_text + f"\n\nπŸ”§ **Direct scheduling call for: {task_description}**\n⏳ *Processing results...*"
485
+
486
+ logger.info(
487
+ f"MCP tool completed with status: {result.get('status', 'unknown')}"
488
+ )
489
+ logger.info(f"MCP result type: {type(result)}")
490
+ logger.info(
491
+ f"MCP result keys: {list(result.keys()) if isinstance(result, dict) else 'Not a dict'}"
492
+ )
493
+
494
+ # Debug the result structure
495
+ if isinstance(result, dict):
496
+ logger.info(f"Result status: {result.get('status')}")
497
+ logger.info(f"Result has schedule: {'schedule' in result}")
498
+ logger.info(
499
+ f"Result has calendar_entries: {'calendar_entries' in result}"
500
+ )
501
+ if "schedule" in result:
502
+ logger.info(
503
+ f"Schedule length: {len(result.get('schedule', []))}"
504
+ )
505
+ if "calendar_entries" in result:
506
+ logger.info(
507
+ f"Calendar entries length: {len(result.get('calendar_entries', []))}"
508
+ )
509
+
510
+ # Check multiple possible success conditions
511
+ is_success = False
512
+ success_msg = ""
513
+
514
+ if result.get("status") == "success":
515
+ is_success = True
516
+ success_msg = "Status is 'success'"
517
+ elif isinstance(result, dict) and "schedule" in result:
518
+ is_success = True
519
+ success_msg = "Result contains schedule data"
520
+ elif isinstance(result, dict) and len(result) > 0:
521
+ is_success = True
522
+ success_msg = "Result contains data"
523
+
524
+ logger.info(f"Success check: {is_success} ({success_msg})")
525
+
526
+ if is_success:
527
+ logger.info(
528
+ "SUCCESS CONDITION MET - Processing successful result"
529
+ )
530
+ schedule = result.get("schedule", [])
531
+ calendar_entries = result.get("calendar_entries", [])
532
+
533
+ # Format the schedule as a table
534
+ if schedule:
535
+ # Create table header
536
+ table_md = "\n\n## πŸ“… **Generated Schedule**\n\n"
537
+ table_md += "| Start Time | End Time | Task | Project | Employee | Duration | Skill | Status |\n"
538
+ table_md += "|------------|----------|------|---------|----------|----------|-------|--------|\n"
539
+
540
+ # Add table rows
541
+ for item in schedule:
542
+ # Use the correct field names from schedule_to_dataframe
543
+ start_time = item.get(
544
+ "Start", item.get("start_time", "N/A")
545
+ )
546
+ end_time = item.get(
547
+ "End", item.get("end_time", "N/A")
548
+ )
549
+ task_name = item.get(
550
+ "Task",
551
+ item.get(
552
+ "task_name", item.get("description", "N/A")
553
+ ),
554
+ )
555
+ project = item.get(
556
+ "Project", item.get("project", "N/A")
557
+ )
558
+ employee = item.get(
559
+ "Employee", item.get("employee", "N/A")
560
+ )
561
+ duration = item.get(
562
+ "Duration (hours)", item.get("duration", "N/A")
563
+ )
564
+ skill = item.get(
565
+ "Required Skill", item.get("skill", "N/A")
566
+ )
567
+
568
+ # Status indicators based on flags
569
+ status_flags = []
570
+ if item.get("Pinned", False):
571
+ status_flags.append("πŸ“Œ Pinned")
572
+ if item.get("Unavailable", False):
573
+ status_flags.append("⚠️ Unavailable")
574
+ if item.get("Undesired", False):
575
+ status_flags.append("😐 Undesired")
576
+ if item.get("Desired", False):
577
+ status_flags.append("βœ… Desired")
578
+
579
+ status = (
580
+ " ".join(status_flags)
581
+ if status_flags
582
+ else "βšͺ Normal"
583
+ )
584
+
585
+ # Format dates/times if they are datetime strings
586
+ if isinstance(start_time, str) and "T" in str(
587
+ start_time
588
+ ):
589
+ try:
590
+ from datetime import datetime
591
+
592
+ dt = datetime.fromisoformat(
593
+ str(start_time).replace("Z", "+00:00")
594
+ )
595
+ start_time = dt.strftime("%m/%d %H:%M")
596
+ except:
597
+ pass
598
+
599
+ if isinstance(end_time, str) and "T" in str(
600
+ end_time
601
+ ):
602
+ try:
603
+ from datetime import datetime
604
+
605
+ dt = datetime.fromisoformat(
606
+ str(end_time).replace("Z", "+00:00")
607
+ )
608
+ end_time = dt.strftime("%m/%d %H:%M")
609
+ except:
610
+ pass
611
+
612
+ # Truncate long task names for table display
613
+ if len(str(task_name)) > 35:
614
+ task_name = str(task_name)[:32] + "..."
615
+
616
+ # Format duration
617
+ if isinstance(duration, (int, float)):
618
+ duration = f"{duration}h"
619
+
620
+ table_md += f"| {start_time} | {end_time} | {task_name} | {project} | {employee} | {duration} | {skill} | {status} |\n"
621
+
622
+ table_md += f"\n**Summary:**\n"
623
+ table_md += f"- πŸ“Š **Total Items:** {len(schedule)}\n"
624
+ table_md += f"- πŸ“… **Calendar Events:** {len(calendar_entries)}\n"
625
+ table_md += f"- βœ… **Status:** Successfully scheduled around existing commitments\n"
626
+
627
+ # Count different types of tasks
628
+ pinned_count = sum(
629
+ 1 for item in schedule if item.get("Pinned", False)
630
+ )
631
+ project_tasks = sum(
632
+ 1
633
+ for item in schedule
634
+ if item.get("Project") == "PROJECT"
635
+ )
636
+ existing_events = sum(
637
+ 1
638
+ for item in schedule
639
+ if item.get("Project") == "EXISTING"
640
+ )
641
+
642
+ table_md += f"- πŸ“Œ **Pinned Events:** {pinned_count}\n"
643
+ table_md += f"- πŸ†• **New Tasks:** {project_tasks}\n"
644
+ table_md += (
645
+ f"- πŸ“… **Existing Events:** {existing_events}\n"
646
+ )
647
+
648
+ # Add JSON data section for debugging
649
+ table_md += f"\n\n<details>\n<summary>πŸ“‹ **Raw JSON Data** (click to expand)</summary>\n\n"
650
+ table_md += "```json\n"
651
+ table_md += safe_json_dumps(result)
652
+ table_md += "\n```\n</details>\n"
653
+
654
+ tool_response = table_md
655
+ else:
656
+ tool_response = f"""
657
+
658
+ πŸ“… **Schedule Generated Successfully!**
659
+
660
+ **Task:** {task_description}
661
+ **Calendar Events Processed:** {len(calendar_entries)}
662
+ **Total Scheduled Items:** {len(schedule)}
663
+
664
+ ⚠️ **No schedule items to display** - This may indicate the task was completed or no scheduling was needed.
665
+
666
+ **Raw Result:**
667
+ ```json
668
+ {safe_json_dumps(result, indent=2)[:1000]}
669
+ ```
670
+ """
671
+
672
+ response_text += tool_response
673
+ logger.info("Added success message with table to response")
674
+ yield response_text
675
+ else:
676
+ logger.warning(f"SUCCESS CONDITION NOT MET")
677
+ error_msg = result.get(
678
+ "error",
679
+ f"Unknown error - result: {safe_json_dumps(result)[:200]}",
680
+ )
681
+ tool_response = f"\n\n❌ **Scheduling Error:** {error_msg}"
682
+ response_text += tool_response
683
+ logger.info("Added error message to response")
684
+ yield response_text
685
+
686
+ except Exception as e:
687
+ logger.error(f"Direct scheduling call failed: {e}")
688
+ import traceback
689
+
690
+ logger.error(f"Full traceback: {traceback.format_exc()}")
691
+ tool_response = f"\n\n❌ **Scheduling failed:** {str(e)}"
692
+ response_text += tool_response
693
+ logger.info("Added exception message to response")
694
+ yield response_text
695
+
696
+ # Always yield final response
697
+ logger.info(f"Final yield: response length {len(response_text)}")
698
+ yield response_text
699
+
700
+ except Exception as e:
701
+ logger.error(f"Error in chat response: {e}")
702
+ import traceback
703
+
704
+ logger.error(f"Full traceback: {traceback.format_exc()}")
705
+ yield f"Error: {str(e)}"
src/ui/pages/mcp_info.py ADDED
@@ -0,0 +1,136 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+
3
+ from utils.logging_config import setup_logging, get_logger
4
+
5
+ # Initialize logging
6
+ setup_logging()
7
+ logger = get_logger(__name__)
8
+
9
+
10
+ def draw_info_page(debug: bool = False):
11
+ with gr.Tab("πŸ“‹ Use as MCP Tool"):
12
+
13
+ gr.Markdown(
14
+ """
15
+ ## πŸ”Œ **Using as MCP Tool**
16
+
17
+ You can use Yuga Planner as an MCP server to integrate scheduling into your AI workflows.
18
+ """
19
+ )
20
+
21
+ gr.Textbox(
22
+ value=get_server_url(),
23
+ label="🌐 MCP Server Endpoint",
24
+ interactive=False,
25
+ max_lines=1,
26
+ )
27
+
28
+ with gr.Accordion("πŸ“ MCP Setup Instructions", open=True):
29
+ gr.Markdown(
30
+ """
31
+ ### 1. **Cursor Setup Instructions (should work from any MCP client!)**
32
+
33
+ **For Cursor AI Editor:**
34
+ 1. Create or edit your MCP configuration file: `~/.cursor/mcp.json`
35
+ 2. Add the yuga-planner server configuration:
36
+ ```json
37
+ {
38
+ "mcpServers": {
39
+ "yuga-planner": {
40
+ "url": -> "Insert the above endpoint URL here"
41
+ }
42
+ }
43
+ }
44
+ ```
45
+ 3. If you already have other MCP servers, add `yuga-planner` to the existing `mcpServers` object
46
+ 4. Restart Cursor to load the new configuration
47
+ 5. The tool will be available in your chat
48
+
49
+ ### 2. **Usage Example**
50
+ """
51
+ )
52
+
53
+ gr.Textbox(
54
+ value="""use yuga-planner mcp tool
55
+ Task Description: Create a new EC2 instance on AWS
56
+
57
+ [Attach your calendar.ics file to provide existing commitments]
58
+
59
+ Tool Response: Optimized schedule created - EC2 setup task assigned to
60
+ available time slots around your existing meetings
61
+ [Returns JSON response with schedule data]
62
+
63
+ User: show all fields as a table, ordered by start date
64
+
65
+ [Displays formatted schedule table with all tasks and calendar events]""",
66
+ label="πŸ’¬ Cursor Chat Usage Example",
67
+ interactive=False,
68
+ lines=10,
69
+ )
70
+
71
+ gr.Markdown(
72
+ """
73
+ ### 3. **What it does**
74
+
75
+ **Personal Task Scheduling with Calendar Integration:**
76
+
77
+ 1. πŸ“… **Parses your calendar** (.ics file) for existing commitments
78
+ 2. πŸ€– **AI breaks down your task** into actionable subtasks using LLamaIndex + Nebius AI
79
+ 3. ⚑ **Constraint-based optimization** finds optimal time slots around your existing schedule
80
+ 4. πŸ“‹ **Returns complete solved schedule** integrated with your personal calendar events
81
+ 5. πŸ•˜ **Respects business hours** (9:00-18:00) and excludes weekends automatically
82
+ 6. πŸ“Š **JSON response format** - Ask to "show all fields as a table, ordered by start date" for readable formatting
83
+
84
+ **Designed for**: Personal productivity and task planning around existing appointments in Cursor.
85
+ """
86
+ )
87
+
88
+ if debug:
89
+ with gr.Tab("πŸ› Debug Info"):
90
+ gr.Markdown(
91
+ """
92
+ # πŸ› Debug Information
93
+
94
+ **Debug Mode Enabled** - Additional system information and controls available.
95
+ """
96
+ )
97
+
98
+ with gr.Accordion("πŸ”§ **Environment Details**", open=True):
99
+ import os
100
+
101
+ env_info = f"""
102
+ **🐍 Python Environment**
103
+ - Debug Mode: {debug}
104
+ - YUGA_DEBUG: {os.getenv('YUGA_DEBUG', 'Not Set')}
105
+ - Nebius API Key: {'βœ… Set' if os.getenv('NEBIUS_API_KEY') else '❌ Not Set'}
106
+ - Nebius Model: {os.getenv('NEBIUS_MODEL', 'Not Set')}
107
+
108
+ **🌐 Server Information**
109
+ - MCP Endpoint: {get_server_url()}
110
+ - Current Working Directory: {os.getcwd()}
111
+ """
112
+ gr.Markdown(env_info)
113
+
114
+ with gr.Accordion("πŸ“Š **System Status**", open=False):
115
+ gr.Markdown(
116
+ """
117
+ **πŸ”„ Service Status**
118
+ - DataService: βœ… Active
119
+ - ScheduleService: βœ… Active
120
+ - StateService: βœ… Active
121
+ - LoggingService: βœ… Active
122
+ - MockProjectService: βœ… Active
123
+
124
+ **πŸ”Œ Integration Status**
125
+ - MCP Server: βœ… Enabled
126
+ - Gradio API: βœ… Active
127
+ - Real-time Logs: βœ… Streaming
128
+ """
129
+ )
130
+
131
+
132
+ def get_server_url():
133
+ try:
134
+ return gr.get_state().server_url + "/gradio_api/mcp/sse"
135
+ except:
136
+ return "https://blackopsrepl-yuga-planner.hf.space/gradio_api/mcp/sse"
tests/test_json_repair.py ADDED
@@ -0,0 +1,179 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import sys
3
+
4
+
5
+ # Add src to path to import our modules
6
+ sys.path.insert(0, "src")
7
+
8
+ from handlers.tool_call_handler import ToolCallAssembler
9
+
10
+ # Import standardized test utilities
11
+ from tests.test_utils import get_test_logger, create_test_results
12
+
13
+ # Initialize standardized test logger
14
+ logger = get_test_logger(__name__)
15
+
16
+
17
+ def test_actual_streaming_error():
18
+ """Test the JSON repair functionality with the actual streaming error pattern"""
19
+
20
+ logger.start_test("Testing JSON repair functionality with actual streaming error")
21
+
22
+ # This is based on the actual error from the logs at character 787
23
+ # The pattern shows base64 data ending abruptly with a quote and then a duplicate JSON object
24
+ broken_json = """{"task_description":"create ec2 on aws","calendar_file_content":"QkVHSU46VkNBTEVOREFSClZFUlNJT046Mi4wClBST0RJRDotLy9pY2FsLm1hcnVkb3QuY29tLy9pQ2FsIEV2ZW50IE1ha2VyCkNBTFNDQUxFOkdSRUdPUklBTgpCRUdJTjpWVElNRVpPTkUKVFpJRDpBZnJpY2EvTGFnb3MKTEFTVC1NT0RJRklFRDoyMDI0MDQyMlQwNTM0NTBaClRaVVJMOmh0dHBzOi8vd3d3LnR6dXJsLm9yZy96b25laW5mby1vdXRsb29rL0FmcmljYS9MYWdvcwpYLUxJQy1MT0NBVElPTjpBZnJpY2EvTGFnb3MKQkVHSU46U1RBTkRBUkQKVFpOQU1FOldBVApUWk9GRlNFVEZST006KzAxMDAKVFpPRkZTRVRUTzorMDEwMApEVFNUQVJUOjE5NzAwMTAxVDAwMDAwMApFTkQ6U1RBTkRBUkQKRU5EOlZUSU1FWk9ORQpCRUdJTjpWRVZFTlQKRFRTVEFNUDoyMDI1MDYyMFQxMzQxMjBaClVJRDpyZWN1ci1tZWV0aW5nLTJAbW9jawpEVFNUQVJUO1RaSUQ9QWZyaWNhL0xhZ29zOjIwMjUwNjAyVDE1MDAwMApEVEVORDtUWklEPUFmcmljYS9MYWdvczoyMDI1MDYwMlQxNjAwMDAKU1VNTUFSWTpQcm9qZWN0IFJldmlldwpFTkQ6VkVWRU5UCkJFR0lOO"{"task_description":"create ec2 on aws","calendar_file_content":"QkVHSU46VkNBTEVOREFSClZFUlNJT046Mi4wClBST0RJRDotLy9pY2FsLm1hcnVkb3QuY29tLy9pQ2FsIEV2ZW50IE1ha2VyCkNBTFNDQUxFOkdSRUdPUklBTgpCRUdJTjpWVElNRVpPTkUKVFpJRDpBZnJpY2EvTGFnb3MKTEFTVC1NT0RJRklFRDoyMDI0MDQyMlQwNTM0NTBaClRaVVJMOmh0dHBzOi8vd3d3LnR6dXJsLm9yZy96b25laW5mby1vdXRsb29rL0FmcmljYS9MYWdvcwpYLUxJQy1MT0NBVElPTjpBZnJpY2EvTGFnb3MKQkVHSU46U1RBTkRBUkQKVFpOQU1FOldBVApUWk9GRlNFVEZST006KzAxMDAKVFpPRkZTRVRUTzorMDEwMApEVFNUQVJUOjE5NzAwMTAxVDAwMDAwMApFTkQ6U1RBTkRBUkQKRU5EOlZUSU1FWk9ORQpCRUdJTjpWRVZFTlQKRFRTVEFNUDoyMDI1MDYyMFQxMzQxMjBaClVJRDpyZWN1ci1tZWV0aW5nLTFAbW9jawpEVFNUQVJUO1RaSUQ9QWZyaWNhL0xhZ29zOjIwMjUwNjAzVDEwMDAwMApEVEVORDtUWklEPUFmcmljYS9MYWdvczoyMDI1MDYwM1QxMTAwMDAKU1VNTUFSWTpUZWFtIFN5bmMKRU5EOlZFVkVOVApCRUdJTjpWRVZFTlQKRFRTVEFNUDoyMDI1MDYyMFQxMzQxMjBaClVJRDpzaW5nbGUtZXZlbnQtMUBtb2NrCkRUU1RBUlQ7VFpJRD1BZnJpY2EvTGFnb3M6MjAyNTA2MDVUMTQwMDAwCkRURU5EO1RaSUQ9QWZyaWNhL0xhZ29zOjIwMjUwNjA1VDE1MDAwMApTVU1NQVJZOkNsaWVudCBDYWxsCkVORDpWRVZFTlQKQkVHSU46VkVWRU5UCkRUU1RBTVA6MjAyNTA2MjBUMTM0MTIwWgpVSUQ6c2luZ2xlLWV2ZW50LTRAbW9jawpEVFNUQVJUO1RaSUQ9QWZyaWNhL0xhZ29zOjIwMjUwNjE2VDE2MDAwMApEVEVORDtUWklEPUFmcmljYS9MYWdvczoyMDI1MDYxNlQxNzAwMDAKU1VNTUFSWTpXb3Jrc2hvcApFTkQ6VkVWRU5UCkJFR0lOOlZFVkVOVApEVFNUQU1QOjIwMjUwNjIwVDEzNDEyMFoKVUlEOnNpbmdsZS1ldmVudC0zQG1vY2sKRFRTVEFSVDtUWklEPUFmcmljYS9MYWdvczoyMDI1MDcwN1QxMTAwMDAKRFRFTkQ7VFpJRD1BZnJpY2EvTGFnb3M6MjAyNTA3MDdUMTIwMDAwClNVTU1BUlk6UGxhbm5pbmcgU2Vzc2lvbgpFTkQ6VkVWRU5UCkJFR0lOOlZFVkVOVApEVFNUQU1QOjIwMjUwNjIwVDEzNDEyMFoKVUlEOnNpbmdsZS1ldmVudC01QG1vY2sKRFRTVEFSVDtUWklEPUFmcmljYS9MYWdvczoyMDI1MDcyMlQwOTAwMDAKRFRFTkQ7VFpJRD1BZnJpY2EvTGFnb3M6MjAyNTA3MjJUMTAwMDAwClNVTU1BUlk6RGVtbwpFTkQ6VkVWRU5UCkVORDpWQ0FMRU5EQVI="}"""
25
+
26
+ logger.debug(f"Broken JSON length: {len(broken_json)}")
27
+
28
+ # Find the error position (where the duplicate starts)
29
+ error_pos = broken_json.find(';"{"task_description"')
30
+ logger.debug(f"Error position (duplicate JSON start): {error_pos}")
31
+
32
+ # Try to parse the broken JSON first to confirm it fails
33
+ json_parse_failed = False
34
+ try:
35
+ json.loads(broken_json)
36
+ logger.error("❌ UNEXPECTED: Broken JSON parsed successfully!")
37
+ except json.JSONDecodeError as e:
38
+ json_parse_failed = True
39
+ logger.info(f"βœ… Expected JSON error at position {e.pos}: {e}")
40
+ logger.debug(f"Error context: '{broken_json[max(0, e.pos-20):e.pos+20]}'")
41
+
42
+ assert json_parse_failed, "Expected broken JSON to fail parsing"
43
+
44
+ # Test the repair function
45
+ assembler = ToolCallAssembler()
46
+ repaired_json = assembler._attempt_json_repair(broken_json)
47
+
48
+ assert repaired_json is not None, "Repair should return a result"
49
+
50
+ logger.info(f"βœ… Repair attempted, result length: {len(repaired_json)}")
51
+ logger.debug(f"Repaired preview: {repaired_json[:200]}...")
52
+
53
+ # Try to parse the repaired JSON
54
+ parsed = json.loads(repaired_json)
55
+ logger.debug(f"Task description: {parsed.get('task_description', 'MISSING')}")
56
+ logger.debug(
57
+ f"Calendar content length: {len(parsed.get('calendar_file_content', ''))}"
58
+ )
59
+
60
+ # Verify expected fields exist
61
+ assert "task_description" in parsed, "Repaired JSON should have task_description"
62
+ assert (
63
+ "calendar_file_content" in parsed
64
+ ), "Repaired JSON should have calendar_file_content"
65
+ assert (
66
+ parsed["task_description"] == "create ec2 on aws"
67
+ ), "Task description should match expected value"
68
+
69
+ logger.pass_test("Repaired JSON parses correctly")
70
+
71
+
72
+ def test_simpler_corruption():
73
+ """Test a simpler case of JSON corruption for baseline functionality"""
74
+
75
+ logger.start_test("Testing simpler JSON corruption")
76
+
77
+ # Missing closing brace
78
+ simple_broken = (
79
+ '{"task_description":"test task","calendar_file_content":"base64data"'
80
+ )
81
+
82
+ assembler = ToolCallAssembler()
83
+ repaired = assembler._attempt_json_repair(simple_broken)
84
+
85
+ assert repaired is not None, "Simple repair should return a result"
86
+
87
+ # Try to parse the repaired JSON
88
+ parsed = json.loads(repaired)
89
+
90
+ assert (
91
+ "task_description" in parsed
92
+ ), "Simple repair should preserve task_description"
93
+ assert parsed["task_description"] == "test task", "Task description should match"
94
+
95
+ logger.pass_test("Simple repair works correctly")
96
+
97
+
98
+ def test_datetime_serialization():
99
+ """Test our datetime serialization fixes"""
100
+
101
+ logger.start_test("Testing datetime serialization")
102
+
103
+ # Import our safe serialization function
104
+ from ui.pages.chat import safe_json_dumps
105
+ from datetime import datetime
106
+
107
+ test_data = {
108
+ "schedule": [
109
+ {
110
+ "task": "Test Task",
111
+ "start_time": datetime(2025, 6, 23, 10, 0),
112
+ "end_time": datetime(2025, 6, 23, 11, 0),
113
+ }
114
+ ],
115
+ "timestamp": datetime.now(),
116
+ }
117
+
118
+ result = safe_json_dumps(test_data, indent=2)
119
+ logger.debug(f"Sample output: {result[:200]}...")
120
+
121
+ # Verify it's valid JSON
122
+ parsed_back = json.loads(result)
123
+ assert "schedule" in parsed_back, "Serialized result should have schedule"
124
+ assert "timestamp" in parsed_back, "Serialized result should have timestamp"
125
+ assert len(parsed_back["schedule"]) == 1, "Schedule should have one item"
126
+
127
+ logger.pass_test("Datetime serialization works correctly")
128
+
129
+
130
+ def test_gradio_format():
131
+ """Test that we're returning the correct format for Gradio messages"""
132
+
133
+ logger.start_test("Testing Gradio message format")
134
+
135
+ # Simulate a proper messages format
136
+ test_history = [
137
+ {"role": "user", "content": "Hello"},
138
+ {"role": "assistant", "content": "Hi there!"},
139
+ {"role": "user", "content": "Create a schedule"},
140
+ ]
141
+
142
+ # This is what our function should return
143
+ expected_format = [
144
+ {"role": "user", "content": "Hello"},
145
+ {"role": "assistant", "content": "Hi there!"},
146
+ {"role": "user", "content": "Create a schedule"},
147
+ {"role": "assistant", "content": "Schedule created successfully!"},
148
+ ]
149
+
150
+ logger.debug("Expected format is list of dicts with 'role' and 'content' keys")
151
+
152
+ # Validate the format
153
+ for i, msg in enumerate(expected_format):
154
+ assert isinstance(msg, dict), f"Message {i} should be a dict, got {type(msg)}"
155
+ assert "role" in msg, f"Message {i} should have 'role' key"
156
+ assert "content" in msg, f"Message {i} should have 'content' key"
157
+ assert msg["role"] in [
158
+ "user",
159
+ "assistant",
160
+ ], f"Message {i} has invalid role: {msg['role']}"
161
+
162
+ logger.pass_test("Message format is correct for Gradio")
163
+
164
+
165
+ if __name__ == "__main__":
166
+ logger.section("JSON Repair and Chat Functionality Tests")
167
+
168
+ # Create test results tracker
169
+ results = create_test_results(logger)
170
+
171
+ # Run tests using the standardized approach
172
+ results.run_test("streaming_error_repair", test_actual_streaming_error)
173
+ results.run_test("simple_json_repair", test_simpler_corruption)
174
+ results.run_test("datetime_serialization", test_datetime_serialization)
175
+ results.run_test("gradio_message_format", test_gradio_format)
176
+
177
+ # Generate summary and exit with appropriate code
178
+ all_passed = results.summary()
179
+ sys.exit(0 if all_passed else 1)
tests/test_tool_assembly.py ADDED
@@ -0,0 +1,177 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Test script for tool call assembly logic
4
+ """
5
+
6
+ import json
7
+ import sys
8
+ import os
9
+
10
+ # Add src to path so we can import modules
11
+ sys.path.insert(0, "src")
12
+
13
+ from handlers.tool_call_handler import ToolCallAssembler
14
+
15
+ # Import standardized test utilities
16
+ from tests.test_utils import get_test_logger, create_test_results
17
+
18
+ # Initialize standardized test logger
19
+ logger = get_test_logger(__name__)
20
+
21
+
22
+ def test_tool_call_assembly():
23
+ """Test the tool call assembler with sample streaming data"""
24
+
25
+ logger.start_test("Testing Tool Call Assembly Logic")
26
+
27
+ assembler = ToolCallAssembler()
28
+
29
+ # Simulate streaming deltas like we saw in the logs
30
+ sample_deltas = [
31
+ # Initial tool call with ID
32
+ {
33
+ "tool_calls": [
34
+ {
35
+ "index": 0,
36
+ "id": "chatcmpl-tool-ca3c56dcd04049cd8baf9a2cde4205d6",
37
+ "function": {"name": "schedule_tasks_with_calendar"},
38
+ "type": "function",
39
+ }
40
+ ]
41
+ },
42
+ # Arguments coming in chunks
43
+ {"tool_calls": [{"index": 0, "function": {"arguments": '{"task_description'}}]},
44
+ {"tool_calls": [{"index": 0, "function": {"arguments": '":"create an'}}]},
45
+ {
46
+ "tool_calls": [
47
+ {"index": 0, "function": {"arguments": " engaging gradio ui"}}
48
+ ]
49
+ },
50
+ {
51
+ "tool_calls": [
52
+ {
53
+ "index": 0,
54
+ "function": {
55
+ "arguments": ' for yuga","calendar_file_content":"test123"}'
56
+ },
57
+ }
58
+ ]
59
+ },
60
+ ]
61
+
62
+ logger.debug("Processing streaming deltas...")
63
+ for i, delta in enumerate(sample_deltas):
64
+ logger.debug(f" Delta {i+1}: {delta}")
65
+ assembler.process_delta(delta)
66
+
67
+ # Show debug info after each delta
68
+ debug_info = assembler.debug_info()
69
+ logger.debug(
70
+ f" -> Tool calls: {debug_info['total_tool_calls']}, Completed: {debug_info['completed_tool_calls']}"
71
+ )
72
+
73
+ if debug_info["tool_calls_detail"]:
74
+ for idx, detail in debug_info["tool_calls_detail"].items():
75
+ logger.debug(
76
+ f" -> Tool {idx}: {detail['function_name']}, Args valid: {detail['is_json_valid']}"
77
+ )
78
+ logger.debug(f" Args preview: {detail['arguments_preview']}")
79
+
80
+ completed_calls = assembler.get_completed_tool_calls()
81
+ logger.info(f"βœ… Completed tool calls: {len(completed_calls)}")
82
+
83
+ for i, tool_call in enumerate(completed_calls):
84
+ logger.debug(f"Tool Call {i+1}:")
85
+ logger.debug(f" ID: {tool_call['id']}")
86
+ logger.debug(f" Function: {tool_call['function']['name']}")
87
+ logger.debug(f" Arguments: {tool_call['function']['arguments']}")
88
+
89
+ # Try to parse arguments
90
+ try:
91
+ args = json.loads(tool_call["function"]["arguments"])
92
+ logger.debug(f" βœ… JSON Valid: {args}")
93
+
94
+ except json.JSONDecodeError as e:
95
+ logger.error(f" ❌ JSON Invalid: {e}")
96
+ raise AssertionError(f"Tool call {i+1} has invalid JSON: {e}")
97
+
98
+ # Verify we got expected results
99
+ assert len(completed_calls) > 0, "Should have at least one completed tool call"
100
+
101
+ logger.pass_test("Normal tool call assembly works correctly")
102
+
103
+
104
+ def test_broken_json():
105
+ """Test with broken JSON to see how we handle it"""
106
+
107
+ logger.start_test("Testing Broken JSON Handling")
108
+
109
+ assembler = ToolCallAssembler()
110
+
111
+ # Simulate incomplete/broken JSON
112
+ broken_deltas = [
113
+ {
114
+ "tool_calls": [
115
+ {
116
+ "index": 0,
117
+ "id": "test-broken",
118
+ "function": {"name": "schedule_tasks_with_calendar"},
119
+ "type": "function",
120
+ }
121
+ ]
122
+ },
123
+ {
124
+ "tool_calls": [
125
+ {"index": 0, "function": {"arguments": '{"task_description":"test'}}
126
+ ]
127
+ },
128
+ # Missing closing quote and brace - should be invalid JSON
129
+ ]
130
+
131
+ for delta in broken_deltas:
132
+ assembler.process_delta(delta)
133
+
134
+ completed_calls = assembler.get_completed_tool_calls()
135
+ debug_info = assembler.debug_info()
136
+
137
+ logger.debug(f"Broken JSON Test Results:")
138
+ logger.debug(f" Total tool calls: {debug_info['total_tool_calls']}")
139
+ logger.debug(f" Completed (valid JSON): {len(completed_calls)}")
140
+ logger.debug(f" Expected: 0 completed (due to broken JSON)")
141
+
142
+ for idx, detail in debug_info["tool_calls_detail"].items():
143
+ logger.debug(f" Tool {idx}: JSON valid = {detail['is_json_valid']}")
144
+ logger.debug(f" Args: {detail['arguments_preview']}")
145
+
146
+ # Should be 0 due to invalid JSON
147
+ expected_completed = 0
148
+ assert (
149
+ len(completed_calls) == expected_completed
150
+ ), f"Expected {expected_completed} completed calls due to broken JSON, got {len(completed_calls)}"
151
+
152
+ logger.pass_test("Broken JSON handling works correctly")
153
+
154
+
155
+ if __name__ == "__main__":
156
+ logger.section("Tool Call Assembly Test Suite")
157
+ logger.info("Testing the isolated tool call assembly logic...")
158
+
159
+ # Create test results tracker
160
+ results = create_test_results(logger)
161
+
162
+ # Run tests using the standardized approach
163
+ results.run_test("normal_assembly", test_tool_call_assembly)
164
+ results.run_test("broken_json_handling", test_broken_json)
165
+
166
+ # Generate summary and exit with appropriate code
167
+ all_passed = results.summary()
168
+
169
+ if all_passed:
170
+ logger.info(
171
+ "πŸŽ‰ All tests passed! Tool call assembly logic is working correctly."
172
+ )
173
+
174
+ else:
175
+ logger.error("πŸ’₯ Some tests failed. Check the logic above.")
176
+
177
+ sys.exit(0 if all_passed else 1)