move backend source to src

backend/src/inference.py

@@ -0,0 +1,60 @@
+import threading
+from stream import Stream
+
+LLAMA3_MODEL = "meta/meta-llama-3-8b-instruct"
+VALID_ROLES = ['user', 'agent', 'system']
+
+class LlamaMessage:
+    def __init__(self, role, content):
+        if role not in VALID_ROLES:
+            raise Exception("message role must be one of: " + ",".join(VALID_ROLES))
+        self.role = role
+        self.content = content
+
+
+class Models:
+    llama_8b = "meta/meta-llama-3-8b-instruct"
+    llama_70b = "meta/meta-llama-3-70b-instruct"
+
+async def replicate_llama_completion(app, messages, model=None):
+    if model is None:
+        model = Models.llama_8b
+
+    # Llama instruct
+    full_prompt = "<|begin_of_text|>"
+    for message in messages:
+        header_name = message.role if message.role != "agent" else "assistant"
+        full_prompt += f"<|start_header_id|>{header_name}<|end_header_id|>\n\n{message.content}<|eot_id|>"
+    full_prompt += f"<|start_header_id|>assistant<|end_header_id|>"
+    
+    content_stream = Stream(app)
+    thread = threading.Thread(
+        target=run_replicate_inference_loop,
+        kwargs={
+            'replicate_client': app.replicate_client,
+            'model': model,
+            'prompt': full_prompt,
+            'content_stream': content_stream,
+        }
+    )
+    thread.start()
+    response = await content_stream.content_wait()
+    return response.lstrip()
+
+
+def run_replicate_inference_loop(replicate_client, model, prompt, content_stream):
+    print("START INFERENCE")
+    completion = replicate_client.run(
+        model,
+        input={
+            "prompt": prompt,
+            "temperature": 0.75,
+            "max_new_tokens": 500,
+            "system_prompt": "" # We handle system prompt in the formatted prompt
+        }
+    )
+    for event in completion:
+        if event:
+            content_stream.feed(str(event))
+    print("END INFERENCE")
+    content_stream.close()