123 lines
4.9 KiB
Python
123 lines
4.9 KiB
Python
|
from google import genai
|
||
|
|
from google.genai import types
|
||
|
|
import os
|
||
|
|
import time
|
||
|
|
from .note_server import NoteServer
|
||
|
|
|
||
|
|
class ObsidianAgent:
|
||
|
|
def __init__(self, api_key: str, server: NoteServer):
|
||
|
|
self.client = genai.Client(api_key=api_key)
|
||
|
|
self.server = server
|
||
|
|
self.model_id = "gemini-3-flash-preview"
|
||
|
|
|
||
|
|
def process_vault(self, input_folder: str, philosophy: str, rewrite_tag: str) -> list[str]:
|
||
|
|
"""
|
||
|
|
Scans the input folder and processes each note using the LLM.
|
||
|
|
Returns a list of action logs.
|
||
|
|
"""
|
||
|
|
logs = []
|
||
|
|
notes = self.server.list_notes(input_folder)
|
||
|
|
|
||
|
|
# Derive Vault Root (Assuming input_folder is inside the vault, e.g. .../Vault/Inbox)
|
||
|
|
# We take the parent directory of the input folder.
|
||
|
|
vault_root = os.path.dirname(input_folder.rstrip(os.sep))
|
||
|
|
|
||
|
|
if not notes:
|
||
|
|
return ["No notes found in input folder."]
|
||
|
|
|
||
|
|
logs.append(f"Found {len(notes)} notes. Processing (Slow mode for Rate Limits)...")
|
||
|
|
|
||
|
|
for i, note_path in enumerate(notes):
|
||
|
|
# RATE LIMITING: Sleep to respect 5 RPM (1 request every 12s + buffer)
|
||
|
|
if i > 0:
|
||
|
|
time.sleep(15)
|
||
|
|
|
||
|
|
content = self.server.read_note(note_path)
|
||
|
|
if not content:
|
||
|
|
continue
|
||
|
|
|
||
|
|
log = self._process_single_note_with_retry(note_path, content, philosophy, rewrite_tag, vault_root)
|
||
|
|
logs.append(log)
|
||
|
|
|
||
|
|
return logs
|
||
|
|
|
||
|
|
def _process_single_note_with_retry(self, note_path: str, content: str, philosophy: str, rewrite_tag: str, vault_root: str) -> str:
|
||
|
|
"""Wraps processing with retry logic for 429 errors."""
|
||
|
|
max_retries = 3
|
||
|
|
for attempt in range(max_retries):
|
||
|
|
try:
|
||
|
|
return self._process_single_note(note_path, content, philosophy, rewrite_tag, vault_root)
|
||
|
|
except Exception as e:
|
||
|
|
if "429" in str(e) or "RESOURCE_EXHAUSTED" in str(e):
|
||
|
|
if attempt < max_retries - 1:
|
||
|
|
wait_time = 60 * (attempt + 1)
|
||
|
|
print(f"Rate limit hit. Waiting {wait_time}s...")
|
||
|
|
time.sleep(wait_time)
|
||
|
|
continue
|
||
|
|
return f"Error processing {os.path.basename(note_path)}: {str(e)}"
|
||
|
|
return f"Failed after retries: {os.path.basename(note_path)}"
|
||
|
|
|
||
|
|
def _process_single_note(self, note_path: str, content: str, philosophy: str, rewrite_tag: str, vault_root: str) -> str:
|
||
|
|
|
||
|
|
# 1. Define Tools
|
||
|
|
def move_note(target_folder: str):
|
||
|
|
"""Moves the note to a folder relative to Vault Root. E.g. 'Science/Biology'."""
|
||
|
|
# Ensure we don't treat target_folder as absolute unless it starts with vault_root
|
||
|
|
if not target_folder.startswith(vault_root):
|
||
|
|
full_target = os.path.join(vault_root, target_folder.strip(os.sep))
|
||
|
|
else:
|
||
|
|
full_target = target_folder
|
||
|
|
|
||
|
|
return self.server.move_note(note_path, full_target)
|
||
|
|
|
||
|
|
def flag_rewrite(reason: str):
|
||
|
|
"""Flags the current note for rewrite by appending a tag and reason."""
|
||
|
|
return self.server.flag_rewrite(note_path, reason, rewrite_tag)
|
||
|
|
|
||
|
|
# 2. Construct Prompt
|
||
|
|
prompt = f"""
|
||
|
|
You are an expert Knowledge Manager for an Obsidian Vault.
|
||
|
|
|
||
|
|
PHILOSOPHY:
|
||
|
|
{philosophy}
|
||
|
|
|
||
|
|
CURRENT NOTE CONTENT:
|
||
|
|
{content[:10000]} # Truncate to avoid massive context if note is huge
|
||
|
|
|
||
|
|
TASK:
|
||
|
|
Analyze the note content against the Philosophy.
|
||
|
|
1. If it fits well into a specific folder in the vault structure, move it there.
|
||
|
|
2. If it is low quality, incomplete, or violates the philosophy, flag it for rewrite.
|
||
|
|
|
||
|
|
You MUST call a tool.
|
||
|
|
"""
|
||
|
|
|
||
|
|
# 3. Call Gemini
|
||
|
|
response = self.client.models.generate_content(
|
||
|
|
model=self.model_id,
|
||
|
|
contents=prompt,
|
||
|
|
config=types.GenerateContentConfig(
|
||
|
|
tools=[move_note, flag_rewrite],
|
||
|
|
tool_config=types.ToolConfig(
|
||
|
|
function_calling_config=types.FunctionCallingConfig(
|
||
|
|
mode="ANY" # Force the model to use a tool
|
||
|
|
)
|
||
|
|
)
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
# 4. Execute Tool
|
||
|
|
if not response.function_calls:
|
||
|
|
return f"Skipped: {os.path.basename(note_path)} (Model failed to call tool)"
|
||
|
|
|
||
|
|
# Execute the first function call found
|
||
|
|
fc = response.function_calls[0]
|
||
|
|
if fc.name == "move_note":
|
||
|
|
result = move_note(**fc.args)
|
||
|
|
return f"[AI] {result}"
|
||
|
|
elif fc.name == "flag_rewrite":
|
||
|
|
result = flag_rewrite(**fc.args)
|
||
|
|
return f"[AI] {result}"
|
||
|
|
else:
|
||
|
|
return f"Error: Unknown tool called {fc.name}"
|