Spaces:
Paused
Paused
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 +59 -50
- src/app.py +4 -140
- src/domain.py +6 -0
- src/handlers/__init__.py +9 -0
- src/handlers/tool_call_handler.py +454 -0
- src/services/__init__.py +2 -0
- src/services/mcp_client.py +80 -0
- src/ui/pages/chat.py +705 -0
- src/ui/pages/mcp_info.py +136 -0
- tests/test_json_repair.py +179 -0
- tests/test_tool_assembly.py +177 -0
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 |
-
###
|
28 |
|
29 |
-
1.
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
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.
|
95 |
-
**Purpose:**
|
96 |
-
- **Access:**
|
97 |
-
- **Input:**
|
98 |
-
- **
|
99 |
-
-
|
|
|
|
|
|
|
|
|
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
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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
|
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:
|
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 |
-
|
|
|
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)
|