Test User
102 lines
3.3 KiB
Python
102 lines
3.3 KiB
Python
"""
|
|
LLM communication module
|
|
"""
|
|
import os
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
from dotenv import load_dotenv
|
|
from openai import OpenAI
|
|
|
|
|
|
class LLMClient:
|
|
def __init__(self, prompts_dir: Path):
|
|
load_dotenv()
|
|
|
|
self.api_url = os.getenv('OPENAI_API_URL', 'https://api.openai.com/v1')
|
|
self.model = os.getenv('OPENAI_MODEL', 'gpt-4')
|
|
self.api_key = os.getenv('OPENAI_API_KEY', '')
|
|
|
|
self.client = OpenAI(
|
|
base_url=self.api_url,
|
|
api_key=self.api_key
|
|
) if self.api_key else None
|
|
|
|
self.prompts_dir = prompts_dir
|
|
if not prompts_dir.exists():
|
|
prompts_dir.mkdir(parents=True)
|
|
self._create_default_prompt()
|
|
|
|
def _create_default_prompt(self):
|
|
"""Create default system prompt"""
|
|
default_prompt = """You are a metadata extraction assistant for markdown documents.
|
|
|
|
Your task is to analyze the document content and suggest appropriate frontmatter metadata.
|
|
|
|
Guidelines:
|
|
- For tags: prefer selecting from the provided existing tags list when appropriate
|
|
- You may create new tags if they better fit the content
|
|
- New tags must follow the rules: lowercase letters, numbers, and underscores only
|
|
- Be concise and accurate
|
|
- Preserve important existing metadata when present"""
|
|
|
|
with open(self.prompts_dir / 'default.txt', 'w', encoding='utf-8') as f:
|
|
f.write(default_prompt)
|
|
|
|
def load_prompt(self, name: str = 'default') -> str:
|
|
"""Load prompt template"""
|
|
prompt_path = self.prompts_dir / f'{name}.txt'
|
|
|
|
if not prompt_path.exists():
|
|
prompt_path = self.prompts_dir / 'default.txt'
|
|
if not prompt_path.exists():
|
|
self._create_default_prompt()
|
|
|
|
with open(prompt_path, 'r', encoding='utf-8') as f:
|
|
return f.read()
|
|
|
|
def infer_metadata(self, field: str, prompt_name: str, original_fm: dict,
|
|
current_fm: dict, document_body: str,
|
|
all_tags: list, schema: dict) -> Optional[str]:
|
|
"""Use LLM to infer metadata field value"""
|
|
if not self.client:
|
|
return None
|
|
|
|
system_prompt = self.load_prompt(prompt_name)
|
|
|
|
user_message = f"""Document content:
|
|
{document_body[:2000]}
|
|
|
|
Original frontmatter:
|
|
{json.dumps(original_fm, indent=2)}
|
|
|
|
Current frontmatter being built:
|
|
{json.dumps(current_fm, indent=2)}
|
|
|
|
Available tags in repository:
|
|
{', '.join(all_tags[:100])}
|
|
|
|
Please provide appropriate metadata for this document following the schema."""
|
|
|
|
try:
|
|
response = self.client.chat.completions.create(
|
|
model=self.model,
|
|
messages=[
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_message}
|
|
],
|
|
response_format={"type": "json_schema", "json_schema": {
|
|
"name": "frontmatter_response",
|
|
"strict": True,
|
|
"schema": schema
|
|
}},
|
|
temperature=0.3
|
|
)
|
|
|
|
result = json.loads(response.choices[0].message.content)
|
|
return result.get(field)
|
|
|
|
except Exception as e:
|
|
print(f" LLM error: {e}")
|
|
return None
|