1import ollama
2from ollama import GenerateResponse
3from ollama import Options
4
5
[docs]
6class OllamaClientManager:
7 def __init__(self):
[docs]
8 self.client = ollama.Client(host="http://localhost:11434")
9 self._model = "llama3.1:8b"
10 self._context_size = 128_000
11
[docs]
12 def embed(self, content: str):
13 response = self.client.embed(
14 model=self._model,
15 input=content,
16 options=Options(num_ctx=self._context_size),
17 )
18 return response.embeddings
19
[docs]
20 def chat(self, prompt: str) -> str:
21 response: GenerateResponse = self.client.generate(
22 model=self._model,
23 prompt=prompt,
24 options=Options(
25 num_ctx=self._context_size,
26 temperature=0,
27 ),
28 )
29 return response.response