64 lines
2.4 KiB
Python
64 lines
2.4 KiB
Python
# services/llm/ollama_service.py
|
|
"""
|
|
Ollama service implementation
|
|
"""
|
|
import ollama
|
|
import json
|
|
from typing import Dict, List
|
|
from .base import LLMService
|
|
from utils.logger import setup_logger
|
|
|
|
|
|
class OllamaService(LLMService):
|
|
def __init__(self, model: str = "qwen3:latest", max_tokens: int = 4000):
|
|
self.model = model
|
|
# Explicitly set the host to avoid potential DNS/proxy issues with 'localhost'
|
|
self.client = ollama.Client(host="127.0.0.1:11434")
|
|
self.max_tokens = max_tokens
|
|
self.logger = setup_logger("ollama")
|
|
|
|
def generate_text(self, prompt: str) -> str:
|
|
self.logger.info(f"--- PROMPT ---\n{prompt}")
|
|
try:
|
|
options = {"num_predict": self.max_tokens}
|
|
response = self.client.generate(
|
|
model=self.model, prompt=prompt, options=options
|
|
)
|
|
response_content = response["response"]
|
|
self.logger.info(f"--- RESPONSE ---\n{response_content}")
|
|
return response_content
|
|
except Exception as e:
|
|
self.logger.error(f"Error in Ollama API call: {e}")
|
|
print(f"Error in Ollama API call: {e}")
|
|
return None
|
|
|
|
def get_similarity_scores(self, texts_pairs: Dict[str, List[str]]) -> List[float]:
|
|
system_prompt = (
|
|
"Evaluate the semantic similarity between the following table of pairs of texts in json format on a scale from 0 to 1. "
|
|
"Return the similarity scores for every row in JSON format as a list of numbers, without any additional text or formatting."
|
|
)
|
|
|
|
request_payload = json.dumps(texts_pairs)
|
|
prompt = f"{system_prompt}\n\n{request_payload}"
|
|
|
|
try:
|
|
options = {"num_predict": self.max_tokens}
|
|
response = self.client.generate(
|
|
model=self.model, prompt=prompt, options=options
|
|
)
|
|
|
|
try:
|
|
scores = json.loads(response["response"].strip())
|
|
if isinstance(scores, dict) and "similarity_scores" in scores:
|
|
return scores["similarity_scores"]
|
|
elif isinstance(scores, list):
|
|
return scores
|
|
else:
|
|
raise ValueError("Unexpected response format")
|
|
except json.JSONDecodeError:
|
|
raise ValueError("Could not decode response as JSON")
|
|
|
|
except Exception as e:
|
|
print(f"Error in Ollama similarity calculation: {e}")
|
|
return None
|