Added LLM route

Hannes221 · Mar 21, 2024 · f00e882 · f00e882
1 parent 3f0ff62
commit f00e882
Show file tree

Hide file tree

Showing 5 changed files with 213 additions and 1 deletion.
diff --git a/backend/llm/__init__.py b/backend/llm/__init__.py
diff --git a/backend/llm/llm_setup.py b/backend/llm/llm_setup.py
@@ -0,0 +1,26 @@
+import os
+import json
+import predictionguard as pg
+
+# Set your Prediction Guard token as an environmental variable.
+
+os.environ["PREDICTIONGUARD_TOKEN"] = "q1VuOjnffJ3NO2oFN8Q9m8vghYc84ld13jaqdF7E"
+
+def get_response(user_message):
+    # Define the system and user messages.
+    messages = [
+        {
+            "role": "system",
+            "content": "You are a helpful assistant. Your model is hosted by Prediction Guard, a leading AI company."
+        },
+        {
+            "role": "user",
+            "content": user_message
+        }
+    ]
+
+    # Get the response using Prediction Guard.
+    result = pg.Chat.create(model="Neural-Chat-7B", messages=messages)
+
+    # Extract and return the response.
+    return result['messages'][-1]['content']
diff --git a/backend/poetry.lock b/backend/poetry.lock