diff --git a/obsidianki/ai/client.py b/obsidianki/ai/client.py index 3be203c..ed4ceb8 100644 --- a/obsidianki/ai/client.py +++ b/obsidianki/ai/client.py @@ -1,8 +1,11 @@ -import os +from __future__ import annotations +from typing import TYPE_CHECKING +if TYPE_CHECKING: + from litellm import completion + from litellm.types.utils import ModelResponse + +import json from typing import List, Dict, Optional, Union, cast -import litellm -from litellm import completion -from litellm.types.utils import ModelResponse from obsidianki.cli.config import console, CONFIG from obsidianki.cli.utils import process_code_blocks, strip_html @@ -13,9 +16,6 @@ AI_RESULT_SET_SIZE = 20 -# Suppress litellm logging -litellm.suppress_debug_info = True - class FlashcardAI: def __init__(self): model_name = getattr(CONFIG, 'model', 'Claude Sonnet 4.5') @@ -146,6 +146,8 @@ def _call_llm( ) -> Optional[ModelResponse]: """Unified LLM call using litellm""" try: + from litellm import completion + from litellm.types.utils import ModelResponse response = completion( model=self.model, messages=[ @@ -163,6 +165,63 @@ def _call_llm( console.print(f"[red]ERROR:[/red] LLM call failed" + str(e)) return None + def _extract_flashcards_from_response( + self, + response: Optional[ModelResponse], + note: Note, + default_tags: Optional[List[str]] = None + ) -> List[Flashcard]: + """Extract and process flashcards from LLM response + + Args: + response: The LLM response containing tool calls + note: The note to associate with flashcards + default_tags: Optional default tags if card doesn't specify any + + Returns: + List of processed Flashcard objects + """ + if not response: + return [] + + try: + message = response.choices[0].message + if not hasattr(message, 'tool_calls') or not message.tool_calls: + console.print("[yellow]WARNING:[/yellow] No flashcards generated - unexpected response format") + return [] + + tool_call = message.tool_calls[0] + arguments = json.loads(tool_call.function.arguments) + flashcard_dicts = arguments.get('flashcards', []) + + flashcard_objects = [] + for card in flashcard_dicts: + front_original = card.get('front', '') + back_original = card.get('back', '') + + # Process code blocks with syntax highlighting + front_processed = process_code_blocks(front_original, CONFIG.syntax_highlighting) + back_processed = process_code_blocks(back_original, CONFIG.syntax_highlighting) + + # Determine tags priority: card's tags > default_tags > note's tags + tags = card.get('tags') or default_tags or note.tags.copy() + + flashcard = Flashcard( + front=front_processed, + back=back_processed, + note=note, + tags=tags, + front_original=front_original, + back_original=back_original + ) + flashcard_objects.append(flashcard) + + return flashcard_objects + + except Exception as e: + console.print(f"[red]ERROR:[/red] Failed to parse flashcards: {e}") + return [] + def generate_flashcards( self, note: Note, @@ -190,43 +249,7 @@ def generate_flashcards( tool_choice=self._get_tool_choice("create_flashcards") ) - if not response: - return [] - - # Extract flashcards from tool call - try: - message = response.choices[0].message - if hasattr(message, 'tool_calls') and message.tool_calls: - tool_call = message.tool_calls[0] - import json - arguments = json.loads(tool_call.function.arguments) - - flashcard_dicts = arguments['flashcards'] - - flashcard_objects = [] - for card in flashcard_dicts: - front_original = card.get('front', '') - back_original = card.get('back', '') - front_processed = process_code_blocks(front_original, CONFIG.syntax_highlighting) - back_processed = process_code_blocks(back_original, CONFIG.syntax_highlighting) - - flashcard = Flashcard( - front=front_processed, - back=back_processed, - note=note, - tags=card.get('tags', note.tags.copy()), - front_original=front_original, - back_original=back_original - ) - flashcard_objects.append(flashcard) - - return flashcard_objects - except Exception as e: - console.print(f"[red]ERROR:[/red] Failed to parse flashcards: {e}") - return [] - - console.print("[yellow]WARNING:[/yellow] No flashcards generated - unexpected response format") - return [] + return self._extract_flashcards_from_response(response, note) def generate_from_query( self, @@ -252,50 +275,16 @@ def generate_from_query( tool_choice=self._get_tool_choice("create_flashcards") ) - if not response: - return [] - - # Extract flashcards - try: - message = response.choices[0].message - if hasattr(message, 'tool_calls') and message.tool_calls: - tool_call = message.tool_calls[0] - import json - flashcard_dicts = json.loads(tool_call.function.arguments).get("flashcards", []) - - # Create virtual Note object for query-based flashcards - virtual_note = Note( - path="query", - filename=f"Query: {query}", - content=query, - tags=["query-generated"], - size=0 - ) - - flashcard_objects = [] - for card in flashcard_dicts: - front_original = card.get('front', '') - back_original = card.get('back', '') - front_processed = process_code_blocks(front_original, CONFIG.syntax_highlighting) - back_processed = process_code_blocks(back_original, CONFIG.syntax_highlighting) - - flashcard = Flashcard( - front=front_processed, - back=back_processed, - note=virtual_note, - tags=card.get('tags', ["query-generated"]), - front_original=front_original, - back_original=back_original - ) - flashcard_objects.append(flashcard) - - return flashcard_objects - except Exception as e: - console.print(f"[red]ERROR:[/red] Failed to parse flashcards: {e}") - return [] + # Create virtual Note object for query-based flashcards + virtual_note = Note( + path="query", + filename=f"Query: {query}", + content=query, + tags=["query-generated"], + size=0 + ) - console.print("[yellow]WARNING:[/yellow] No flashcards generated - unexpected response format") - return [] + return self._extract_flashcards_from_response(response, virtual_note, default_tags=["query-generated"]) def generate_from_note_query(self, note: Note, query: str, target_cards: int, previous_fronts: List[str] | None = None, deck_examples: List[Dict[str, str]] | None = None) -> List[Flashcard]: """Generate flashcards by extracting specific information from a note based on a query""" @@ -324,41 +313,7 @@ def generate_from_note_query(self, note: Note, query: str, target_cards: int, pr tool_choice=self._get_tool_choice("create_flashcards") ) - if not response: - return [] - - # Extract flashcards - try: - message = response.choices[0].message - if hasattr(message, 'tool_calls') and message.tool_calls: - tool_call = message.tool_calls[0] - import json - flashcard_dicts = json.loads(tool_call.function.arguments).get("flashcards", []) - - flashcard_objects = [] - for card in flashcard_dicts: - front_original = card.get('front', '') - back_original = card.get('back', '') - front_processed = process_code_blocks(front_original, CONFIG.syntax_highlighting) - back_processed = process_code_blocks(back_original, CONFIG.syntax_highlighting) - - flashcard = Flashcard( - front=front_processed, - back=back_processed, - note=note, - tags=card.get('tags', note.tags.copy()), - front_original=front_original, - back_original=back_original - ) - flashcard_objects.append(flashcard) - - return flashcard_objects - except Exception as e: - console.print(f"[red]ERROR:[/red] Failed to parse flashcards: {e}") - return [] - - console.print("[yellow]WARNING:[/yellow] No flashcards generated - unexpected response format") - return [] + return self._extract_flashcards_from_response(response, note) def find_with_agent(self, natural_request: str, sample_size: int | None = None, bias_strength: float | None = None) -> List[Note]: """Use multi-turn agent with tool calling to find notes via iterative DQL refinement""" @@ -413,7 +368,6 @@ def find_with_agent(self, natural_request: str, sample_size: int | None = None, if hasattr(message, 'tool_calls') and message.tool_calls: for tool_call in message.tool_calls: tool_name = tool_call.function.name - import json tool_input = json.loads(tool_call.function.arguments) if tool_name == "execute_dql_query": @@ -602,7 +556,6 @@ def edit_cards(self, cards: List[Dict[str, str]], query: str) -> List[Dict[str, message = response.choices[0].message if hasattr(message, 'tool_calls') and message.tool_calls: tool_call = message.tool_calls[0] - import json flashcard_data = json.loads(tool_call.function.arguments) if "flashcards" in flashcard_data: