my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	from typing import Dict, List, Optional, Union
#2
#3	try:
#4	from ollama import Client
#5	except ImportError:
#6	raise ImportError("The 'ollama' library is required. Please install it using 'pip install ollama'.")
#7
#8	from mem0.configs.llms.base import BaseLlmConfig
#9	from mem0.configs.llms.ollama import OllamaConfig
#10	from mem0.llms.base import LLMBase
#11
#12
#13	class OllamaLLM(LLMBase):
#14	def __init__(self, config: Optional[Union[BaseLlmConfig, OllamaConfig, Dict]] = None):
#15	# Convert to OllamaConfig if needed
#16	if config is None:
#17	config = OllamaConfig()
#18	elif isinstance(config, dict):
#19	config = OllamaConfig(**config)
#20	elif isinstance(config, BaseLlmConfig) and not isinstance(config, OllamaConfig):
#21	# Convert BaseLlmConfig to OllamaConfig
#22	config = OllamaConfig(
#23	model=config.model,
#24	temperature=config.temperature,
#25	api_key=config.api_key,
#26	max_tokens=config.max_tokens,
#27	top_p=config.top_p,
#28	top_k=config.top_k,
#29	enable_vision=config.enable_vision,
#30	vision_details=config.vision_details,
#31	http_client_proxies=config.http_client,
#32	)
#33
#34	super().__init__(config)
#35
#36	if not self.config.model:
#37	self.config.model = "llama3.1:70b"
#38
#39	self.client = Client(host=self.config.ollama_base_url)
#40
#41	def _parse_response(self, response, tools):
#42	"""
#43	Process the response based on whether tools are used or not.
#44
#45	Args:
#46	response: The raw response from API.
#47	tools: The list of tools provided in the request.
#48
#49	Returns:
#50	str or dict: The processed response.
#51	"""
#52	# Get the content from response
#53	if isinstance(response, dict):
#54	content = response["message"]["content"]
#55	else:
#56	content = response.message.content
#57
#58	if tools:
#59	processed_response = {
#60	"content": content,
#61	"tool_calls": [],
#62	}
#63
#64	# Ollama doesn't support tool calls in the same way, so we return the content
#65	return processed_response
#66	else:
#67	return content
#68
#69	def generate_response(
#70	self,
#71	messages: List[Dict[str, str]],
#72	response_format=None,
#73	tools: Optional[List[Dict]] = None,
#74	tool_choice: str = "auto",
#75	**kwargs,
#76	):
#77	"""
#78	Generate a response based on the given messages using Ollama.
#79
#80	Args:
#81	messages (list): List of message dicts containing 'role' and 'content'.
#82	response_format (str or object, optional): Format of the response. Defaults to "text".
#83	tools (list, optional): List of tools that the model can call. Defaults to None.
#84	tool_choice (str, optional): Tool choice method. Defaults to "auto".
#85	**kwargs: Additional Ollama-specific parameters.
#86
#87	Returns:
#88	str: The generated response.
#89	"""
#90	# Build parameters for Ollama
#91	params = {
#92	"model": self.config.model,
#93	"messages": messages,
#94	}
#95
#96	# Handle JSON response format by using Ollama's native format parameter
#97	if response_format and response_format.get("type") == "json_object":
#98	params["format"] = "json"
#99	# Also add JSON format instruction to the last message as a fallback
#100	if messages and messages[-1]["role"] == "user":
#101	messages[-1]["content"] += "\n\nPlease respond with valid JSON only."
#102	else:
#103	messages.append({"role": "user", "content": "Please respond with valid JSON only."})
#104
#105	# Add options for Ollama (temperature, num_predict, top_p)
#106	options = {
#107	"temperature": self.config.temperature,
#108	"num_predict": self.config.max_tokens,
#109	"top_p": self.config.top_p,
#110	}
#111	params["options"] = options
#112
#113	# Remove OpenAI-specific parameters that Ollama doesn't support
#114	params.pop("max_tokens", None) # Ollama uses different parameter names
#115
#116	response = self.client.chat(**params)
#117	return self._parse_response(response, tools)
#118

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public