Skip to content

Commit e60df3d

Browse files
additional param support added
1 parent 43a97fb commit e60df3d

1 file changed

Lines changed: 266 additions & 0 deletions

File tree

examples/additional_params_demo.py

Lines changed: 266 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,266 @@
1+
"""
2+
Additional Params Demo
3+
4+
This example demonstrates how to use the additional_params feature
5+
to attach custom metadata to tasks that will be included in the task response.
6+
7+
Use case: When you need to track additional information with your tasks
8+
like proxy URLs, public URLs, request IDs, tracking info, etc.
9+
"""
10+
11+
import time
12+
import redis
13+
from modelq import ModelQ
14+
from pydantic import BaseModel
15+
16+
17+
# Initialize Redis and ModelQ
18+
redis_client = redis.Redis(host="localhost", port=6379, db=0)
19+
modelq_app = ModelQ(redis_client=redis_client)
20+
21+
22+
# ---------------------------------------------------------------------------
23+
# Example 1: Simple task with additional_params
24+
# ---------------------------------------------------------------------------
25+
26+
@modelq_app.task()
27+
def process_image(image_url: str, width: int, height: int):
28+
"""Simulates image processing."""
29+
print(f"Processing image: {image_url} to size {width}x{height}")
30+
time.sleep(2)
31+
return f"Processed {image_url}"
32+
33+
34+
# Call the task with additional_params
35+
print("Example 1: Image processing with proxy and public links")
36+
print("-" * 60)
37+
38+
task1 = process_image(
39+
"https://example.com/image.jpg",
40+
800,
41+
600,
42+
additional_params={
43+
"proxy_links": [
44+
"http://proxy1.example.com/image_processed.jpg",
45+
"http://proxy2.example.com/image_processed.jpg"
46+
],
47+
"public_links": "http://cdn.example.com/image_processed.jpg",
48+
"request_id": "req_12345",
49+
"user_id": "user_789"
50+
}
51+
)
52+
53+
print(f"Task created: {task1.task_id}")
54+
print(f"Additional params: {task1.additional_params}")
55+
print()
56+
57+
# Get task info immediately (while queued)
58+
task_info = modelq_app.get_task_details(task1.task_id)
59+
print("Task info (queued):")
60+
print(f" Task ID: {task_info['task_id']}")
61+
print(f" Status: {task_info['status']}")
62+
print(f" Proxy Links: {task_info.get('proxy_links', [])}")
63+
print(f" Public Links: {task_info.get('public_links', 'N/A')}")
64+
print(f" Request ID: {task_info.get('request_id', 'N/A')}")
65+
print(f" User ID: {task_info.get('user_id', 'N/A')}")
66+
print()
67+
68+
69+
# ---------------------------------------------------------------------------
70+
# Example 2: Task with Pydantic schema and additional_params
71+
# ---------------------------------------------------------------------------
72+
73+
class TextGenerationInput(BaseModel):
74+
prompt: str
75+
max_tokens: int
76+
temperature: float
77+
78+
79+
class TextGenerationOutput(BaseModel):
80+
text: str
81+
tokens_used: int
82+
83+
84+
@modelq_app.task(schema=TextGenerationInput, returns=TextGenerationOutput)
85+
def generate_text(params: TextGenerationInput):
86+
"""Simulates text generation."""
87+
print(f"Generating text for prompt: {params.prompt[:50]}...")
88+
time.sleep(1)
89+
return TextGenerationOutput(
90+
text=f"Generated text based on: {params.prompt}",
91+
tokens_used=100
92+
)
93+
94+
95+
print("\nExample 2: Text generation with Pydantic schema and metadata")
96+
print("-" * 60)
97+
98+
task2 = generate_text(
99+
TextGenerationInput(
100+
prompt="Write a story about a robot",
101+
max_tokens=500,
102+
temperature=0.7
103+
),
104+
additional_params={
105+
"model_name": "gpt-4",
106+
"api_version": "v1",
107+
"cost_estimate": 0.002,
108+
"billing_id": "billing_xyz",
109+
"tags": ["generation", "story", "robot"]
110+
}
111+
)
112+
113+
print(f"Task created: {task2.task_id}")
114+
print(f"Additional params: {task2.additional_params}")
115+
print()
116+
117+
task2_info = modelq_app.get_task_details(task2.task_id)
118+
print("Task info (queued):")
119+
print(f" Task ID: {task2_info['task_id']}")
120+
print(f" Status: {task2_info['status']}")
121+
print(f" Model: {task2_info.get('model_name', 'N/A')}")
122+
print(f" API Version: {task2_info.get('api_version', 'N/A')}")
123+
print(f" Cost Estimate: ${task2_info.get('cost_estimate', 0)}")
124+
print(f" Tags: {task2_info.get('tags', [])}")
125+
print()
126+
127+
128+
# ---------------------------------------------------------------------------
129+
# Example 3: Task without additional_params (backward compatibility)
130+
# ---------------------------------------------------------------------------
131+
132+
@modelq_app.task()
133+
def simple_calculation(a: int, b: int):
134+
"""Simple task without additional params."""
135+
time.sleep(1)
136+
return a + b
137+
138+
139+
print("\nExample 3: Simple task without additional_params")
140+
print("-" * 60)
141+
142+
task3 = simple_calculation(10, 20)
143+
144+
print(f"Task created: {task3.task_id}")
145+
print(f"Additional params: {task3.additional_params}") # Should be empty dict
146+
print()
147+
148+
task3_info = modelq_app.get_task_details(task3.task_id)
149+
print("Task info (queued):")
150+
print(f" Task ID: {task3_info['task_id']}")
151+
print(f" Status: {task3_info['status']}")
152+
print(f" Has proxy_links: {'proxy_links' in task3_info}")
153+
print(f" Has public_links: {'public_links' in task3_info}")
154+
print()
155+
156+
157+
# ---------------------------------------------------------------------------
158+
# Example 4: Complete workflow with worker
159+
# ---------------------------------------------------------------------------
160+
161+
print("\nExample 4: Complete workflow with worker")
162+
print("-" * 60)
163+
print("Starting worker to process tasks...")
164+
print("(In production, run worker in a separate process)")
165+
print()
166+
167+
# Start worker in background (for demo purposes, run for limited time)
168+
import threading
169+
170+
171+
def run_worker_for_demo():
172+
"""Run worker for a short time to process demo tasks."""
173+
modelq_app.start_workers(no_of_workers=2)
174+
# Worker will process tasks in background
175+
176+
177+
# Start worker thread
178+
worker_thread = threading.Thread(target=run_worker_for_demo, daemon=True)
179+
worker_thread.start()
180+
181+
# Give worker time to start
182+
time.sleep(1)
183+
184+
# Create a task to process
185+
print("Creating img2img task with additional params...")
186+
task4 = process_image(
187+
"https://example.com/input.jpg",
188+
1024,
189+
1024,
190+
additional_params={
191+
"proxy_links": ["http://proxy.cdn.com/output.jpg"],
192+
"public_links": "http://cdn.example.com/output.jpg",
193+
"webhook_url": "https://api.example.com/webhook",
194+
"priority": "high",
195+
"user_metadata": {
196+
"user_id": "user_123",
197+
"session_id": "sess_456"
198+
}
199+
}
200+
)
201+
202+
print(f"Task ID: {task4.task_id}")
203+
print()
204+
205+
# Wait for task to complete
206+
print("Waiting for task to complete...")
207+
try:
208+
result = task4.get_result(redis_client, timeout=10)
209+
print(f"Result: {result}")
210+
print()
211+
212+
# Get final task details
213+
final_info = modelq_app.get_task_details(task4.task_id)
214+
print("Final task info:")
215+
print(f" Status: {final_info['status']}")
216+
print(f" Result: {final_info.get('result', 'N/A')}")
217+
print(f" Proxy Links: {final_info.get('proxy_links', [])}")
218+
print(f" Public Links: {final_info.get('public_links', 'N/A')}")
219+
print(f" Webhook URL: {final_info.get('webhook_url', 'N/A')}")
220+
print(f" Priority: {final_info.get('priority', 'N/A')}")
221+
print(f" User Metadata: {final_info.get('user_metadata', {})}")
222+
print()
223+
224+
except Exception as e:
225+
print(f"Error: {e}")
226+
227+
228+
# ---------------------------------------------------------------------------
229+
# Example 5: API response format
230+
# ---------------------------------------------------------------------------
231+
232+
print("\nExample 5: Simulated API response format")
233+
print("-" * 60)
234+
print("This shows how your task response would look in an API:")
235+
print()
236+
237+
# Create a task
238+
api_task = process_image(
239+
"https://example.com/api_image.jpg",
240+
512,
241+
512,
242+
additional_params={
243+
"proxy_links": [
244+
"http://proxy1.example.com/result.jpg",
245+
"http://proxy2.example.com/result.jpg"
246+
],
247+
"public_links": "http://cdn.example.com/result.jpg"
248+
}
249+
)
250+
251+
# Simulate API response
252+
import json
253+
254+
api_response = api_task.to_dict()
255+
print(json.dumps(api_response, indent=2))
256+
print()
257+
258+
print("=" * 60)
259+
print("Demo complete!")
260+
print()
261+
print("Key takeaways:")
262+
print("1. Use additional_params to attach custom metadata to tasks")
263+
print("2. Additional params appear at the root level of task responses")
264+
print("3. Works with both regular tasks and Pydantic-validated tasks")
265+
print("4. Completely optional - tasks work normally without it")
266+
print("5. Perfect for tracking proxy URLs, public URLs, request IDs, etc.")

0 commit comments

Comments
 (0)