omni-cortex 1.1.0__tar.gz → 1.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (62) hide show
  1. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/PKG-INFO +1 -1
  2. omni_cortex-1.2.0/dashboard/backend/image_service.py +533 -0
  3. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/dashboard/backend/main.py +119 -1
  4. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/dashboard/backend/models.py +45 -0
  5. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/pyproject.toml +1 -1
  6. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/.gitignore +0 -0
  7. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/LICENSE +0 -0
  8. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/README.md +0 -0
  9. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/dashboard/backend/chat_service.py +0 -0
  10. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/dashboard/backend/database.py +0 -0
  11. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/dashboard/backend/logging_config.py +0 -0
  12. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/dashboard/backend/project_config.py +0 -0
  13. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/dashboard/backend/project_scanner.py +0 -0
  14. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/dashboard/backend/pyproject.toml +0 -0
  15. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/dashboard/backend/uv.lock +0 -0
  16. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/dashboard/backend/websocket_manager.py +0 -0
  17. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/hooks/post_tool_use.py +0 -0
  18. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/hooks/pre_tool_use.py +0 -0
  19. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/hooks/stop.py +0 -0
  20. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/hooks/subagent_stop.py +0 -0
  21. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/__init__.py +0 -0
  22. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/categorization/__init__.py +0 -0
  23. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/categorization/auto_tags.py +0 -0
  24. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/categorization/auto_type.py +0 -0
  25. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/config.py +0 -0
  26. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/dashboard.py +0 -0
  27. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/database/__init__.py +0 -0
  28. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/database/connection.py +0 -0
  29. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/database/migrations.py +0 -0
  30. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/database/schema.py +0 -0
  31. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/database/sync.py +0 -0
  32. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/decay/__init__.py +0 -0
  33. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/decay/importance.py +0 -0
  34. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/embeddings/__init__.py +0 -0
  35. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/embeddings/local.py +0 -0
  36. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/models/__init__.py +0 -0
  37. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/models/activity.py +0 -0
  38. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/models/agent.py +0 -0
  39. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/models/memory.py +0 -0
  40. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/models/relationship.py +0 -0
  41. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/models/session.py +0 -0
  42. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/resources/__init__.py +0 -0
  43. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/search/__init__.py +0 -0
  44. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/search/hybrid.py +0 -0
  45. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/search/keyword.py +0 -0
  46. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/search/ranking.py +0 -0
  47. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/search/semantic.py +0 -0
  48. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/server.py +0 -0
  49. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/setup.py +0 -0
  50. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/tools/__init__.py +0 -0
  51. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/tools/activities.py +0 -0
  52. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/tools/memories.py +0 -0
  53. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/tools/sessions.py +0 -0
  54. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/tools/utilities.py +0 -0
  55. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/utils/__init__.py +0 -0
  56. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/utils/formatting.py +0 -0
  57. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/utils/ids.py +0 -0
  58. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/utils/timestamps.py +0 -0
  59. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/omni_cortex/utils/truncation.py +0 -0
  60. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/scripts/import_ken_memories.py +0 -0
  61. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/scripts/populate_session_data.py +0 -0
  62. {omni_cortex-1.1.0 → omni_cortex-1.2.0}/scripts/setup.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: omni-cortex
3
- Version: 1.1.0
3
+ Version: 1.2.0
4
4
  Summary: Give Claude Code a perfect memory - auto-logs everything, searches smartly, and gets smarter over time
5
5
  Project-URL: Homepage, https://github.com/AllCytes/Omni-Cortex
6
6
  Project-URL: Repository, https://github.com/AllCytes/Omni-Cortex
@@ -0,0 +1,533 @@
1
+ """Image generation service using Nano Banana Pro (gemini-3-pro-image-preview)."""
2
+
3
+ import base64
4
+ import os
5
+ import uuid
6
+ from dataclasses import dataclass, field
7
+ from enum import Enum
8
+ from typing import Optional
9
+
10
+ from dotenv import load_dotenv
11
+
12
+ from database import get_memory_by_id
13
+
14
+ load_dotenv()
15
+
16
+
17
+ class ImagePreset(str, Enum):
18
+ """Preset templates for common image types."""
19
+ INFOGRAPHIC = "infographic"
20
+ KEY_INSIGHTS = "key_insights"
21
+ TIPS_TRICKS = "tips_tricks"
22
+ QUOTE_CARD = "quote_card"
23
+ WORKFLOW = "workflow"
24
+ COMPARISON = "comparison"
25
+ SUMMARY_CARD = "summary_card"
26
+ CUSTOM = "custom"
27
+
28
+
29
+ # Preset system prompts
30
+ PRESET_PROMPTS = {
31
+ ImagePreset.INFOGRAPHIC: """Create a professional infographic with:
32
+ - Clear visual hierarchy with icons and sections
33
+ - Bold header/title at top
34
+ - 3-5 key points with visual elements
35
+ - Clean, modern design with good use of whitespace
36
+ - Professional color scheme""",
37
+
38
+ ImagePreset.KEY_INSIGHTS: """Create a clean insights card showing:
39
+ - "Key Insights" or similar header
40
+ - 3-5 bullet points with key takeaways
41
+ - Each insight is concise (1-2 lines max)
42
+ - Clean typography, easy to read
43
+ - Subtle design elements""",
44
+
45
+ ImagePreset.TIPS_TRICKS: """Create a tips card showing:
46
+ - Numbered tips (1, 2, 3, etc.) with icons
47
+ - Each tip is actionable and clear
48
+ - Visual styling that's engaging
49
+ - Good contrast and readability""",
50
+
51
+ ImagePreset.QUOTE_CARD: """Create a quote card with:
52
+ - The key quote in large, styled text
53
+ - Attribution below the quote
54
+ - Elegant, minimalist design
55
+ - Suitable for social media sharing""",
56
+
57
+ ImagePreset.WORKFLOW: """Create a workflow diagram showing:
58
+ - Step-by-step process with arrows/connectors
59
+ - Each step clearly labeled
60
+ - Visual flow from start to finish
61
+ - Professional diagrammatic style""",
62
+
63
+ ImagePreset.COMPARISON: """Create a comparison visual showing:
64
+ - Side-by-side or pros/cons layout
65
+ - Clear distinction between options
66
+ - Visual indicators (checkmarks, icons)
67
+ - Balanced, professional presentation""",
68
+
69
+ ImagePreset.SUMMARY_CARD: """Create a summary card with:
70
+ - Brief title/header
71
+ - Key stats or metrics highlighted
72
+ - Concise overview text
73
+ - Clean, scannable layout""",
74
+
75
+ ImagePreset.CUSTOM: "" # User provides full prompt
76
+ }
77
+
78
+ # Default aspect ratios for presets
79
+ PRESET_ASPECT_RATIOS = {
80
+ ImagePreset.INFOGRAPHIC: "9:16",
81
+ ImagePreset.KEY_INSIGHTS: "1:1",
82
+ ImagePreset.TIPS_TRICKS: "4:5",
83
+ ImagePreset.QUOTE_CARD: "1:1",
84
+ ImagePreset.WORKFLOW: "16:9",
85
+ ImagePreset.COMPARISON: "16:9",
86
+ ImagePreset.SUMMARY_CARD: "4:3",
87
+ ImagePreset.CUSTOM: "16:9",
88
+ }
89
+
90
+
91
+ @dataclass
92
+ class SingleImageRequest:
93
+ """Request for a single image within a batch."""
94
+ preset: ImagePreset = ImagePreset.CUSTOM
95
+ custom_prompt: str = ""
96
+ aspect_ratio: str = "16:9"
97
+ image_size: str = "2K"
98
+
99
+
100
+ @dataclass
101
+ class ImageGenerationResult:
102
+ """Result for a single generated image."""
103
+ success: bool
104
+ image_data: Optional[str] = None # Base64 encoded
105
+ mime_type: str = "image/png"
106
+ text_response: Optional[str] = None
107
+ thought_signature: Optional[str] = None
108
+ error: Optional[str] = None
109
+ index: int = 0 # Position in batch
110
+ image_id: Optional[str] = None
111
+
112
+
113
+ @dataclass
114
+ class BatchImageResult:
115
+ """Result for batch image generation."""
116
+ success: bool
117
+ images: list[ImageGenerationResult] = field(default_factory=list)
118
+ errors: list[str] = field(default_factory=list)
119
+
120
+
121
+ @dataclass
122
+ class ConversationTurn:
123
+ role: str # "user" or "model"
124
+ text: Optional[str] = None
125
+ image_data: Optional[str] = None
126
+ thought_signature: Optional[str] = None
127
+
128
+
129
+ class ImageGenerationService:
130
+ def __init__(self):
131
+ self._api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
132
+ self._client = None
133
+ # Per-image conversation history for multi-turn editing
134
+ self._image_conversations: dict[str, list[ConversationTurn]] = {}
135
+
136
+ def _get_client(self):
137
+ """Get or create the Gemini client."""
138
+ if self._client is None and self._api_key:
139
+ try:
140
+ from google import genai
141
+ self._client = genai.Client(api_key=self._api_key)
142
+ except ImportError:
143
+ return None
144
+ return self._client
145
+
146
+ def is_available(self) -> bool:
147
+ """Check if image generation service is available."""
148
+ if not self._api_key:
149
+ return False
150
+ try:
151
+ from google import genai
152
+ return True
153
+ except ImportError:
154
+ return False
155
+
156
+ def build_memory_context(self, db_path: str, memory_ids: list[str]) -> str:
157
+ """Build context string from selected memories."""
158
+ memories = []
159
+ for mem_id in memory_ids:
160
+ memory = get_memory_by_id(db_path, mem_id)
161
+ if memory:
162
+ memories.append(f"""
163
+ Memory: {memory.memory_type}
164
+ Content: {memory.content}
165
+ Context: {memory.context or 'N/A'}
166
+ Tags: {', '.join(memory.tags) if memory.tags else 'N/A'}
167
+ """)
168
+ return "\n---\n".join(memories)
169
+
170
+ def build_chat_context(self, chat_messages: list[dict]) -> str:
171
+ """Build context string from recent chat conversation."""
172
+ if not chat_messages:
173
+ return ""
174
+
175
+ context_parts = ["Recent conversation context:"]
176
+ for msg in chat_messages[-10:]: # Last 10 messages
177
+ role = msg.get("role", "user")
178
+ content = msg.get("content", "")
179
+ context_parts.append(f"{role}: {content}")
180
+
181
+ return "\n".join(context_parts)
182
+
183
+ def _build_prompt_with_preset(
184
+ self,
185
+ request: SingleImageRequest,
186
+ memory_context: str,
187
+ chat_context: str
188
+ ) -> str:
189
+ """Build full prompt combining preset, custom prompt, and context."""
190
+ parts = []
191
+
192
+ # Add memory context
193
+ if memory_context:
194
+ parts.append(f"Based on the following memories:\n\n{memory_context}")
195
+
196
+ # Add chat context
197
+ if chat_context:
198
+ parts.append(f"\n{chat_context}")
199
+
200
+ # Add preset prompt (if not custom)
201
+ if request.preset != ImagePreset.CUSTOM:
202
+ preset_prompt = PRESET_PROMPTS.get(request.preset, "")
203
+ if preset_prompt:
204
+ parts.append(f"\nImage style guidance:\n{preset_prompt}")
205
+
206
+ # Add user's custom prompt
207
+ if request.custom_prompt:
208
+ parts.append(f"\nUser request: {request.custom_prompt}")
209
+
210
+ parts.append("\nGenerate a professional, high-quality image optimized for social media sharing.")
211
+
212
+ return "\n".join(parts)
213
+
214
+ async def generate_single_image(
215
+ self,
216
+ request: SingleImageRequest,
217
+ memory_context: str,
218
+ chat_context: str = "",
219
+ conversation_history: list[dict] = None,
220
+ use_search_grounding: bool = False,
221
+ image_id: str = None,
222
+ ) -> ImageGenerationResult:
223
+ """Generate a single image based on request and context."""
224
+ client = self._get_client()
225
+ if not client:
226
+ return ImageGenerationResult(
227
+ success=False,
228
+ error="API key not configured or google-genai not installed"
229
+ )
230
+
231
+ try:
232
+ from google.genai import types
233
+ except ImportError:
234
+ return ImageGenerationResult(
235
+ success=False,
236
+ error="google-genai package not installed"
237
+ )
238
+
239
+ # Generate image ID if not provided
240
+ if not image_id:
241
+ image_id = f"img_{uuid.uuid4().hex[:8]}"
242
+
243
+ # Build the full prompt
244
+ full_prompt = self._build_prompt_with_preset(
245
+ request, memory_context, chat_context
246
+ )
247
+
248
+ # Build contents with conversation history for multi-turn editing
249
+ contents = []
250
+
251
+ # Use image-specific conversation history if editing
252
+ if image_id and image_id in self._image_conversations:
253
+ for turn in self._image_conversations[image_id]:
254
+ parts = []
255
+ if turn.text:
256
+ part = {"text": turn.text}
257
+ if turn.thought_signature:
258
+ part["thoughtSignature"] = turn.thought_signature
259
+ parts.append(part)
260
+ if turn.image_data:
261
+ part = {
262
+ "inlineData": {
263
+ "mimeType": "image/png",
264
+ "data": turn.image_data
265
+ }
266
+ }
267
+ if turn.thought_signature:
268
+ part["thoughtSignature"] = turn.thought_signature
269
+ parts.append(part)
270
+ contents.append({
271
+ "role": turn.role,
272
+ "parts": parts
273
+ })
274
+ elif conversation_history:
275
+ # Use provided conversation history
276
+ for turn in conversation_history:
277
+ parts = []
278
+ if turn.get("text"):
279
+ part = {"text": turn["text"]}
280
+ if turn.get("thought_signature"):
281
+ part["thoughtSignature"] = turn["thought_signature"]
282
+ parts.append(part)
283
+ if turn.get("image_data"):
284
+ part = {
285
+ "inlineData": {
286
+ "mimeType": "image/png",
287
+ "data": turn["image_data"]
288
+ }
289
+ }
290
+ if turn.get("thought_signature"):
291
+ part["thoughtSignature"] = turn["thought_signature"]
292
+ parts.append(part)
293
+ contents.append({
294
+ "role": turn["role"],
295
+ "parts": parts
296
+ })
297
+
298
+ # Add current prompt
299
+ contents.append({
300
+ "role": "user",
301
+ "parts": [{"text": full_prompt}]
302
+ })
303
+
304
+ # Configure image settings
305
+ config = types.GenerateContentConfig(
306
+ response_modalities=["IMAGE", "TEXT"],
307
+ )
308
+
309
+ if use_search_grounding:
310
+ config.tools = [{"google_search": {}}]
311
+
312
+ try:
313
+ response = client.models.generate_content(
314
+ model="gemini-2.0-flash-preview-image-generation",
315
+ contents=contents,
316
+ config=config
317
+ )
318
+
319
+ # Extract image and thought signatures
320
+ image_data = None
321
+ text_response = None
322
+ thought_signature = None
323
+
324
+ if response.candidates and response.candidates[0].content:
325
+ for part in response.candidates[0].content.parts:
326
+ if hasattr(part, 'inline_data') and part.inline_data:
327
+ image_data = base64.b64encode(part.inline_data.data).decode()
328
+ if hasattr(part, 'text') and part.text:
329
+ text_response = part.text
330
+ if hasattr(part, 'thought_signature') and part.thought_signature:
331
+ thought_signature = part.thought_signature
332
+
333
+ # Store conversation for this image (for editing)
334
+ if image_id and image_data:
335
+ if image_id not in self._image_conversations:
336
+ self._image_conversations[image_id] = []
337
+ self._image_conversations[image_id].append(
338
+ ConversationTurn(role="user", text=full_prompt)
339
+ )
340
+ self._image_conversations[image_id].append(
341
+ ConversationTurn(
342
+ role="model",
343
+ text=text_response,
344
+ image_data=image_data,
345
+ thought_signature=thought_signature
346
+ )
347
+ )
348
+
349
+ return ImageGenerationResult(
350
+ success=image_data is not None,
351
+ image_data=image_data,
352
+ text_response=text_response,
353
+ thought_signature=thought_signature,
354
+ image_id=image_id,
355
+ error=None if image_data else "No image generated"
356
+ )
357
+
358
+ except Exception as e:
359
+ return ImageGenerationResult(
360
+ success=False,
361
+ error=str(e),
362
+ image_id=image_id
363
+ )
364
+
365
+ async def generate_batch(
366
+ self,
367
+ requests: list[SingleImageRequest],
368
+ memory_context: str,
369
+ chat_context: str = "",
370
+ use_search_grounding: bool = False,
371
+ ) -> BatchImageResult:
372
+ """Generate multiple images with different settings."""
373
+ results = []
374
+ errors = []
375
+
376
+ for i, request in enumerate(requests):
377
+ # Generate unique ID for each image in batch
378
+ image_id = f"batch_{uuid.uuid4().hex[:8]}_{i}"
379
+
380
+ result = await self.generate_single_image(
381
+ request=request,
382
+ memory_context=memory_context,
383
+ chat_context=chat_context,
384
+ use_search_grounding=use_search_grounding,
385
+ image_id=image_id
386
+ )
387
+ result.index = i
388
+ results.append(result)
389
+
390
+ if not result.success:
391
+ errors.append(f"Image {i+1}: {result.error}")
392
+
393
+ return BatchImageResult(
394
+ success=len(errors) == 0,
395
+ images=results,
396
+ errors=errors
397
+ )
398
+
399
+ async def refine_image(
400
+ self,
401
+ image_id: str,
402
+ refinement_prompt: str,
403
+ aspect_ratio: str = None,
404
+ image_size: str = None
405
+ ) -> ImageGenerationResult:
406
+ """Refine an existing image using its conversation history."""
407
+ client = self._get_client()
408
+ if not client:
409
+ return ImageGenerationResult(
410
+ success=False,
411
+ error="API key not configured"
412
+ )
413
+
414
+ if image_id not in self._image_conversations:
415
+ return ImageGenerationResult(
416
+ success=False,
417
+ error="No conversation history found for this image"
418
+ )
419
+
420
+ try:
421
+ from google.genai import types
422
+ except ImportError:
423
+ return ImageGenerationResult(
424
+ success=False,
425
+ error="google-genai package not installed"
426
+ )
427
+
428
+ # Build contents from conversation history
429
+ contents = []
430
+
431
+ for turn in self._image_conversations[image_id]:
432
+ parts = []
433
+ if turn.text:
434
+ part = {"text": turn.text}
435
+ if turn.thought_signature:
436
+ part["thoughtSignature"] = turn.thought_signature
437
+ parts.append(part)
438
+ if turn.image_data:
439
+ part = {
440
+ "inlineData": {
441
+ "mimeType": "image/png",
442
+ "data": turn.image_data
443
+ }
444
+ }
445
+ if turn.thought_signature:
446
+ part["thoughtSignature"] = turn.thought_signature
447
+ parts.append(part)
448
+ contents.append({
449
+ "role": turn.role,
450
+ "parts": parts
451
+ })
452
+
453
+ # Add refinement prompt
454
+ contents.append({
455
+ "role": "user",
456
+ "parts": [{"text": refinement_prompt}]
457
+ })
458
+
459
+ # Configure - use defaults or provided values
460
+ config = types.GenerateContentConfig(
461
+ response_modalities=["IMAGE", "TEXT"],
462
+ )
463
+
464
+ try:
465
+ response = client.models.generate_content(
466
+ model="gemini-2.0-flash-preview-image-generation",
467
+ contents=contents,
468
+ config=config
469
+ )
470
+
471
+ image_data = None
472
+ text_response = None
473
+ thought_signature = None
474
+
475
+ if response.candidates and response.candidates[0].content:
476
+ for part in response.candidates[0].content.parts:
477
+ if hasattr(part, 'inline_data') and part.inline_data:
478
+ image_data = base64.b64encode(part.inline_data.data).decode()
479
+ if hasattr(part, 'text') and part.text:
480
+ text_response = part.text
481
+ if hasattr(part, 'thought_signature') and part.thought_signature:
482
+ thought_signature = part.thought_signature
483
+
484
+ # Update conversation history
485
+ self._image_conversations[image_id].append(
486
+ ConversationTurn(role="user", text=refinement_prompt)
487
+ )
488
+ self._image_conversations[image_id].append(
489
+ ConversationTurn(
490
+ role="model",
491
+ text=text_response,
492
+ image_data=image_data,
493
+ thought_signature=thought_signature
494
+ )
495
+ )
496
+
497
+ return ImageGenerationResult(
498
+ success=image_data is not None,
499
+ image_data=image_data,
500
+ text_response=text_response,
501
+ thought_signature=thought_signature,
502
+ image_id=image_id,
503
+ error=None if image_data else "No image generated"
504
+ )
505
+
506
+ except Exception as e:
507
+ return ImageGenerationResult(
508
+ success=False,
509
+ error=str(e),
510
+ image_id=image_id
511
+ )
512
+
513
+ def clear_conversation(self, image_id: str = None):
514
+ """Clear conversation history. If image_id provided, clear only that image."""
515
+ if image_id:
516
+ self._image_conversations.pop(image_id, None)
517
+ else:
518
+ self._image_conversations.clear()
519
+
520
+ def get_presets(self) -> list[dict]:
521
+ """Get available presets with their default settings."""
522
+ return [
523
+ {
524
+ "value": preset.value,
525
+ "label": preset.value.replace("_", " ").title(),
526
+ "default_aspect": PRESET_ASPECT_RATIOS.get(preset, "16:9")
527
+ }
528
+ for preset in ImagePreset
529
+ ]
530
+
531
+
532
+ # Singleton instance
533
+ image_service = ImageGenerationService()
@@ -39,7 +39,21 @@ from database import (
39
39
  update_memory,
40
40
  )
41
41
  from logging_config import log_success, log_error
42
- from models import ChatRequest, ChatResponse, ConversationSaveRequest, ConversationSaveResponse, FilterParams, MemoryUpdate, ProjectInfo, ProjectRegistration
42
+ from models import (
43
+ ChatRequest,
44
+ ChatResponse,
45
+ ConversationSaveRequest,
46
+ ConversationSaveResponse,
47
+ FilterParams,
48
+ MemoryUpdate,
49
+ ProjectInfo,
50
+ ProjectRegistration,
51
+ BatchImageGenerationRequest,
52
+ BatchImageGenerationResponse,
53
+ ImageRefineRequest,
54
+ SingleImageRequestModel,
55
+ SingleImageResponseModel,
56
+ )
43
57
  from project_config import (
44
58
  load_config,
45
59
  add_registered_project,
@@ -51,6 +65,7 @@ from project_config import (
51
65
  from project_scanner import scan_projects
52
66
  from websocket_manager import manager
53
67
  import chat_service
68
+ from image_service import image_service, ImagePreset, SingleImageRequest
54
69
 
55
70
 
56
71
  class DatabaseChangeHandler(FileSystemEventHandler):
@@ -630,6 +645,109 @@ async def save_chat_conversation(
630
645
  raise
631
646
 
632
647
 
648
+ # --- Image Generation Endpoints ---
649
+
650
+
651
+ @app.get("/api/image/status")
652
+ async def get_image_status():
653
+ """Check if image generation is available."""
654
+ return {
655
+ "available": image_service.is_available(),
656
+ "message": "Image generation ready" if image_service.is_available()
657
+ else "Configure GEMINI_API_KEY and install google-genai for image generation",
658
+ }
659
+
660
+
661
+ @app.get("/api/image/presets")
662
+ async def get_image_presets():
663
+ """Get available image preset templates."""
664
+ return {"presets": image_service.get_presets()}
665
+
666
+
667
+ @app.post("/api/image/generate-batch", response_model=BatchImageGenerationResponse)
668
+ async def generate_images_batch(
669
+ request: BatchImageGenerationRequest,
670
+ db_path: str = Query(..., alias="project", description="Path to the database file"),
671
+ ):
672
+ """Generate multiple images with different presets/prompts."""
673
+ # Validate image count
674
+ if len(request.images) not in [1, 2, 4]:
675
+ return BatchImageGenerationResponse(
676
+ success=False,
677
+ errors=["Must request 1, 2, or 4 images"]
678
+ )
679
+
680
+ # Build memory context
681
+ memory_context = ""
682
+ if request.memory_ids:
683
+ memory_context = image_service.build_memory_context(db_path, request.memory_ids)
684
+
685
+ # Build chat context
686
+ chat_context = image_service.build_chat_context(request.chat_messages)
687
+
688
+ # Convert request models to internal format
689
+ image_requests = [
690
+ SingleImageRequest(
691
+ preset=ImagePreset(img.preset),
692
+ custom_prompt=img.custom_prompt,
693
+ aspect_ratio=img.aspect_ratio,
694
+ image_size=img.image_size
695
+ )
696
+ for img in request.images
697
+ ]
698
+
699
+ result = await image_service.generate_batch(
700
+ requests=image_requests,
701
+ memory_context=memory_context,
702
+ chat_context=chat_context,
703
+ use_search_grounding=request.use_search_grounding
704
+ )
705
+
706
+ return BatchImageGenerationResponse(
707
+ success=result.success,
708
+ images=[
709
+ SingleImageResponseModel(
710
+ success=img.success,
711
+ image_data=img.image_data,
712
+ text_response=img.text_response,
713
+ thought_signature=img.thought_signature,
714
+ image_id=img.image_id,
715
+ error=img.error,
716
+ index=img.index
717
+ )
718
+ for img in result.images
719
+ ],
720
+ errors=result.errors
721
+ )
722
+
723
+
724
+ @app.post("/api/image/refine", response_model=SingleImageResponseModel)
725
+ async def refine_image(request: ImageRefineRequest):
726
+ """Refine an existing generated image with a new prompt."""
727
+ result = await image_service.refine_image(
728
+ image_id=request.image_id,
729
+ refinement_prompt=request.refinement_prompt,
730
+ aspect_ratio=request.aspect_ratio,
731
+ image_size=request.image_size
732
+ )
733
+
734
+ return SingleImageResponseModel(
735
+ success=result.success,
736
+ image_data=result.image_data,
737
+ text_response=result.text_response,
738
+ thought_signature=result.thought_signature,
739
+ image_id=result.image_id,
740
+ error=result.error
741
+ )
742
+
743
+
744
+ @app.post("/api/image/clear-conversation")
745
+ async def clear_image_conversation(image_id: Optional[str] = None):
746
+ """Clear image conversation history. If image_id provided, clear only that image."""
747
+ image_service.clear_conversation(image_id)
748
+ return {"status": "cleared", "image_id": image_id}
749
+
750
+
633
751
  # --- WebSocket Endpoint ---
634
752
 
635
753
 
@@ -186,3 +186,48 @@ class ConversationSaveResponse(BaseModel):
186
186
 
187
187
  memory_id: str
188
188
  summary: str
189
+
190
+
191
+ # --- Image Generation Models ---
192
+
193
+
194
+ class SingleImageRequestModel(BaseModel):
195
+ """Request for a single image in a batch."""
196
+ preset: str = "custom" # Maps to ImagePreset enum
197
+ custom_prompt: str = ""
198
+ aspect_ratio: str = "16:9"
199
+ image_size: str = "2K"
200
+
201
+
202
+ class BatchImageGenerationRequest(BaseModel):
203
+ """Request for generating multiple images."""
204
+ images: list[SingleImageRequestModel] # 1, 2, or 4 images
205
+ memory_ids: list[str] = []
206
+ chat_messages: list[dict] = [] # Recent chat for context
207
+ use_search_grounding: bool = False
208
+
209
+
210
+ class ImageRefineRequest(BaseModel):
211
+ """Request for refining an existing image."""
212
+ image_id: str
213
+ refinement_prompt: str
214
+ aspect_ratio: Optional[str] = None
215
+ image_size: Optional[str] = None
216
+
217
+
218
+ class SingleImageResponseModel(BaseModel):
219
+ """Response for a single generated image."""
220
+ success: bool
221
+ image_data: Optional[str] = None # Base64 encoded
222
+ text_response: Optional[str] = None
223
+ thought_signature: Optional[str] = None
224
+ image_id: Optional[str] = None
225
+ error: Optional[str] = None
226
+ index: int = 0
227
+
228
+
229
+ class BatchImageGenerationResponse(BaseModel):
230
+ """Response for batch image generation."""
231
+ success: bool
232
+ images: list[SingleImageResponseModel] = []
233
+ errors: list[str] = []
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "omni-cortex"
7
- version = "1.1.0"
7
+ version = "1.2.0"
8
8
  description = "Give Claude Code a perfect memory - auto-logs everything, searches smartly, and gets smarter over time"
9
9
  readme = "README.md"
10
10
  license = "MIT"
File without changes
File without changes
File without changes
File without changes