Update app.py

app.py

@@ -1,7 +1,33 @@
 from fastapi import FastAPI
+import os
 
+from custom_llm import CustomLLM
+
+from pydantic import BaseModel
+from langchain.prompts import PromptTemplate
+
+class ConversationPost(BaseModel):
+    tenant: str | None = None
+    module: str | None = None
+    question: str
+
+
+API_TOKEN = os.environ['HF_API_KEY']
 app = FastAPI()
+prompt = PromptTemplate.from_template("""<|im_start|>system
+Kamu adalah Asisten AI yang dikembangkan oleh Jonthan Jordan. Jawablah pertanyaan user secara ketat dalam Bahasa Indonesia<|im_end|>
+<|im_start|>user
+{question}<|im_end|>
+<|im_start|>assistant
+""")
+llm = prompt | CustomLLM(repo_id="Qwen/Qwen-VL-Chat", model_type='text-generation', api_token=API_TOKEN, max_new_tokens=150).bind(stop=['<|im_end|>'])
+
 
 @app.get("/")
 def greet_json():
-    return {"Hello": "World!"}
+    return {"Hello": "World!"}
+
+
+@app.post("/conversation"):
+async def conversation(data : ConversationPost):
+    return {"output":llm.invoke({"question":data.question})}