55 lines
2.1 KiB
Python
55 lines
2.1 KiB
Python
import requests
|
|
import json
|
|
from typing import Dict, Optional
|
|
|
|
class LLMClient:
|
|
def __init__(self, endpoint: str = 'http://192.168.1.74:1234', model: str = 'local'):
|
|
self.endpoint = endpoint
|
|
self.model = model
|
|
self.local_ollama = 'http://localhost:11434'
|
|
|
|
def summarize(self, text: str, max_length: int = 200) -> Dict:
|
|
prompt = f"Summarize the following in {max_length} chars or less:\n\n{text[:2000]}"
|
|
return self._query(prompt)
|
|
|
|
def extract_topics(self, text: str) -> Dict:
|
|
prompt = f"Extract 5-10 key topics/tags from this text. Return as comma-separated list:\n\n{text[:2000]}"
|
|
return self._query(prompt)
|
|
|
|
def classify_content(self, text: str) -> Dict:
|
|
prompt = f"Classify this content. Return: category, topics, has_pii (yes/no), quality (high/medium/low):\n\n{text[:1000]}"
|
|
return self._query(prompt)
|
|
|
|
def _query(self, prompt: str, use_local: bool = False) -> Dict:
|
|
try:
|
|
endpoint = self.local_ollama if use_local else self.endpoint
|
|
|
|
if use_local:
|
|
response = requests.post(
|
|
f'{endpoint}/api/generate',
|
|
json={'model': 'llama3.2', 'prompt': prompt, 'stream': False},
|
|
timeout=30
|
|
)
|
|
else:
|
|
response = requests.post(
|
|
f'{endpoint}/v1/chat/completions',
|
|
json={
|
|
'model': self.model,
|
|
'messages': [{'role': 'user', 'content': prompt}],
|
|
'max_tokens': 500
|
|
},
|
|
timeout=30
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
data = response.json()
|
|
if use_local:
|
|
return {'success': True, 'text': data.get('response', '')}
|
|
else:
|
|
return {'success': True, 'text': data['choices'][0]['message']['content']}
|
|
else:
|
|
return {'success': False, 'error': f'HTTP {response.status_code}'}
|
|
|
|
except Exception as e:
|
|
return {'success': False, 'error': str(e)}
|