File size: 2,910 Bytes
e3278e4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
"""
Ollama /chat/completion calls handled in llm_http_handler.py
[TODO]: migrate embeddings to a base handler as well.
"""
import asyncio
from typing import Any, Dict, List
import litellm
from litellm.types.utils import EmbeddingResponse
# ollama wants plain base64 jpeg/png files as images. strip any leading dataURI
# and convert to jpeg if necessary.
async def ollama_aembeddings(
api_base: str,
model: str,
prompts: List[str],
model_response: EmbeddingResponse,
optional_params: dict,
logging_obj: Any,
encoding: Any,
):
if api_base.endswith("/api/embed"):
url = api_base
else:
url = f"{api_base}/api/embed"
## Load Config
config = litellm.OllamaConfig.get_config()
for k, v in config.items():
if (
k not in optional_params
): # completion(top_k=3) > cohere_config(top_k=3) <- allows for dynamic variables to be passed in
optional_params[k] = v
data: Dict[str, Any] = {"model": model, "input": prompts}
special_optional_params = ["truncate", "options", "keep_alive"]
for k, v in optional_params.items():
if k in special_optional_params:
data[k] = v
else:
# Ensure "options" is a dictionary before updating it
data.setdefault("options", {})
if isinstance(data["options"], dict):
data["options"].update({k: v})
total_input_tokens = 0
output_data = []
response = await litellm.module_level_aclient.post(url=url, json=data)
response_json = response.json()
embeddings: List[List[float]] = response_json["embeddings"]
for idx, emb in enumerate(embeddings):
output_data.append({"object": "embedding", "index": idx, "embedding": emb})
input_tokens = response_json.get("prompt_eval_count") or len(
encoding.encode("".join(prompt for prompt in prompts))
)
total_input_tokens += input_tokens
model_response.object = "list"
model_response.data = output_data
model_response.model = "ollama/" + model
setattr(
model_response,
"usage",
litellm.Usage(
prompt_tokens=total_input_tokens,
completion_tokens=total_input_tokens,
total_tokens=total_input_tokens,
prompt_tokens_details=None,
completion_tokens_details=None,
),
)
return model_response
def ollama_embeddings(
api_base: str,
model: str,
prompts: list,
optional_params: dict,
model_response: EmbeddingResponse,
logging_obj: Any,
encoding=None,
):
return asyncio.run(
ollama_aembeddings(
api_base=api_base,
model=model,
prompts=prompts,
model_response=model_response,
optional_params=optional_params,
logging_obj=logging_obj,
encoding=encoding,
)
)
|