add linter and test ollama in python

2024-04-29 15:31:01 +07:00
parent b2866d073c
commit c23b1d42bf
6 changed files with 984 additions and 294 deletions
@@ -0,0 +1,58 @@
+import asyncio
+import re
+
+import ollama
+
+
+async def chat(answers: dict[int, str], request: str = 'музычку пожалуйста') -> str:
+    answers_str = ''.join([f'{key}. {value}\n' for key, value in answers.items()])
+    messages = [
+        {
+            'role': 'system',
+            'content': 'Ты должен ответить на вопрос пользователя числом или набором чисел.'
+                       'Пример твоего ответа: 1.'
+                       'Ещё один пример: 4 и 3'
+                       ' Вот список:\n' + answers_str
+        },
+        {
+            'role': 'user',
+            'content': request
+        }
+    ]
+    async_client = ollama.AsyncClient()
+    response = await async_client.chat(
+        model='llama3:8b',
+        messages=messages,
+        options={
+            'temperature': 0.5,
+        }
+    )
+    return response.get("message").get("content")
+
+
+async def main():
+    answers = {
+        1: "включить музыку",
+        2: "остановить музыку",
+        3: "следующая композиция",
+        4: "предыдущая композиция",
+        5: 'погода',
+        99: "ничего из этого"
+    }
+    number = None
+    while not number:
+        response_content = await chat(answers, "останови музыку и скажи погоду")
+        print(response_content)
+        number = [
+            int(i)
+            for i in re.findall(r'\d+', response_content)
+            if int(i) in answers.keys()
+        ]
+        if not number:
+            print("Не удалось определить номер. Пожалуйста, попробуйте снова.")
+    print(f"Выбранный номер: {number}")
+
+
+if __name__ == '__main__':
+    ollama.pull("llama3:8b")
+    asyncio.run(main())