projects/pyrag3/vector_service.py
2026-04-05 17:30:07 -07:00

72 lines
2.5 KiB
Python

import requests
import logging
import math
logger = logging.getLogger(__name__)
class VectorService:
def __init__(self):
# Hardcoded for portability in standalone binary
self.api_key = "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"
self.model = "nvidia/llama-nemotron-embed-vl-1b-v2:free"
self.url = "https://openrouter.ai/api/v1/embeddings"
self.headers = {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"
}
def get_embedding(self, text=None, image_b64=None, mime_type="image/jpeg"):
"""Fetch multimodal embedding from OpenRouter API."""
if not text and not image_b64: return None
# Structure payload for multimodal embedding
content = []
if text:
content.append({"type": "text", "text": text[:10000]}) # Truncate large text
if image_b64:
content.append({
"type": "image_url",
"image_url": {"url": f"data:{mime_type};base64,{image_b64}"}
})
payload = {
"model": self.model,
"input": [{"content": content}]
}
try:
resp = requests.post(self.url, headers=self.headers, json=payload, timeout=45)
if resp.status_code == 200:
data = resp.json()
if "data" in data and len(data["data"]) > 0:
return data["data"][0]["embedding"]
else:
logger.error(f"OpenRouter Error {resp.status_code}: {resp.text}")
except Exception as e:
logger.error(f"Failed to get embedding: {e}")
return None
@staticmethod
def cosine_similarity(vec1, vec2):
"""Pure Python implementation of Cosine Similarity."""
if not vec1 or not vec2 or len(vec1) != len(vec2):
return 0.0
dot_product = sum(a * b for a, b in zip(vec1, vec2))
magnitude1 = math.sqrt(sum(a * a for a in vec1))
magnitude2 = math.sqrt(sum(b * b for b in vec2))
if magnitude1 == 0 or magnitude2 == 0:
return 0.0
return dot_product / (magnitude1 * magnitude2)
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
vs = VectorService()
emb = vs.get_embedding("Hello world")
if emb:
print(f"Embedding length: {len(emb)}")
sim = vs.cosine_similarity(emb, emb)
print(f"Self-similarity: {sim}")