my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import os
#2	from typing import Dict, List, Optional
#3
#4	import requests
#5
#6	from mem0.configs.llms.base import BaseLlmConfig
#7	from mem0.llms.base import LLMBase
#8
#9
#10	class SarvamLLM(LLMBase):
#11	def __init__(self, config: Optional[BaseLlmConfig] = None):
#12	super().__init__(config)
#13
#14	# Set default model if not provided
#15	if not self.config.model:
#16	self.config.model = "sarvam-m"
#17
#18	# Get API key from config or environment variable
#19	self.api_key = self.config.api_key or os.getenv("SARVAM_API_KEY")
#20
#21	if not self.api_key:
#22	raise ValueError(
#23	"Sarvam API key is required. Set SARVAM_API_KEY environment variable or provide api_key in config."
#24	)
#25
#26	# Set base URL - use config value or environment or default
#27	self.base_url = (
#28	getattr(self.config, "sarvam_base_url", None) or os.getenv("SARVAM_API_BASE") or "https://api.sarvam.ai/v1"
#29	)
#30
#31	def generate_response(self, messages: List[Dict[str, str]], response_format=None) -> str:
#32	"""
#33	Generate a response based on the given messages using Sarvam-M.
#34
#35	Args:
#36	messages (list): List of message dicts containing 'role' and 'content'.
#37	response_format (str or object, optional): Format of the response.
#38	Currently not used by Sarvam API.
#39
#40	Returns:
#41	str: The generated response.
#42	"""
#43	url = f"{self.base_url}/chat/completions"
#44
#45	headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
#46
#47	# Prepare the request payload
#48	params = {
#49	"messages": messages,
#50	"model": self.config.model if isinstance(self.config.model, str) else "sarvam-m",
#51	}
#52
#53	# Add standard parameters that already exist in BaseLlmConfig
#54	if self.config.temperature is not None:
#55	params["temperature"] = self.config.temperature
#56
#57	if self.config.max_tokens is not None:
#58	params["max_tokens"] = self.config.max_tokens
#59
#60	if self.config.top_p is not None:
#61	params["top_p"] = self.config.top_p
#62
#63	# Handle Sarvam-specific parameters if model is passed as dict
#64	if isinstance(self.config.model, dict):
#65	# Extract model name
#66	params["model"] = self.config.model.get("name", "sarvam-m")
#67
#68	# Add Sarvam-specific parameters
#69	sarvam_specific_params = ["reasoning_effort", "frequency_penalty", "presence_penalty", "seed", "stop", "n"]
#70
#71	for param in sarvam_specific_params:
#72	if param in self.config.model:
#73	params[param] = self.config.model[param]
#74
#75	try:
#76	response = requests.post(url, headers=headers, json=params, timeout=30)
#77	response.raise_for_status()
#78
#79	result = response.json()
#80
#81	if "choices" in result and len(result["choices"]) > 0:
#82	return result["choices"][0]["message"]["content"]
#83	else:
#84	raise ValueError("No response choices found in Sarvam API response")
#85
#86	except requests.exceptions.RequestException as e:
#87	raise RuntimeError(f"Sarvam API request failed: {e}")
#88	except KeyError as e:
#89	raise ValueError(f"Unexpected response format from Sarvam API: {e}")
#90

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public