diff --git a/py/api.py b/py/api.py index e5ffba9..09373b5 100644 --- a/py/api.py +++ b/py/api.py @@ -1,7 +1,29 @@ import requests +import json + class api: + def initialize_ollama(self, prompt): + ollama_url = "http://localhost:11434" + model = "phi3.5" - def send_message(message): - answer = "Test" - return answer \ No newline at end of file + response = requests.post( + f"{ollama_url}/api/generate", json={"model": model, "prompt": prompt} + ) + + if response.status_code == 200: + response_data = [] + for line in response.iter_lines(): + line_decoded = line.decode("utf-8") + line_raw = json.loads(line_decoded) + response_data.append(line_raw["response"]) + + final_response = "".join(response_data) + generated_text = final_response.splitlines()[-1] + return generated_text + else: + return "Error: " + response.text + + def send_message(self, message): + answer = self.initialize_ollama(message) + return answer