Update demo project.py

demo project.py

@@ -1,6 +1,5 @@
 import gradio as gr
 from langchain_groq import ChatGroq
-import os
 from langgraph.graph import StateGraph, START, END
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_chroma import Chroma
@@ -12,6 +11,10 @@ import time
 import os
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_huggingface import HuggingFaceEmbeddings
+from bytez import Bytez
+
+BYTEZ_API_KEY = os.getenv("BYTEZ_API_KEY")
+sdk = Bytez(BYTEZ_API_KEY)
 
 class State(TypedDict):
     query: str
@@ -25,23 +28,24 @@ class checker_class(BaseModel):
     is_relevant: bool = Field(description="Check whether the given query is relevant to the company.")
 
 def invoke_llm(query):
-    llm = ChatGroq(model='llama-3.3-70b-versatile')
+    llm = sdk.model("openai/gpt-4o-mini")
+    Inp = [{"role": "user", "content": query}]
     try:
-        res = llm.invoke(query)
+        res = llm.run(Inp)
     except:
         time.sleep(60)
         res = llm.invoke(query)
-    return res.content
+    return res[0]['content']
     
 def invoke_relevance_checker_llm(query):
-    llm = ChatGroq(model='llama-3.3-70b-versatile')
-    checker_llm = llm.with_structured_output(checker_class)
+    llm = sdk.model("openai/gpt-4o-mini")
+    Inp = [{"role": "user", "content": query}]
     try:
-        res = checker_llm.invoke([HumanMessage(content=query)])
+        res = llm.run(Inp)
     except:
         time.sleep(60)
-        res = checker_llm.invoke([HumanMessage(content=query)])
-    return res.is_relevant
+        res = llm.invoke(query)
+    return res[0]['content']
 
 def safety_checker(state:State):
     llm = ChatGroq(model='meta-llama/llama-guard-4-12b')
@@ -53,7 +57,7 @@ def safety_checker(state:State):
         return {'is_safe':False, 'answer':"<SAFETY CHECKER> That prompt was harmful, please try something else"}
         
 def relevance_checker(state:State):
-    prompt = "You are a lenient relevance-checking assistant. You will be given a user query and a company description. Your job is to decide whether the query is relevant to the company.\n✅ Approve most queries that are even loosely related.\n🚫 Only reject queries that are **clearly unrelated** or have **no connection at all**.\n\n"
+    prompt = "You are a lenient relevance-checking assistant. You will be given a user query and a company description. Your job is to decide whether the query is relevant to the company.\n✅ Approve most queries that are even loosely related.\n🚫 Only reject queries that are **clearly unrelated** or have **no connection at all**. Return only 'True' if it is relevant, otherwise 'False'\n\n"
     prompt += f"\nQuery: {state['query']}"
     prompt += f"\nDescription: {state['company_description']}"
     res = invoke_relevance_checker_llm(prompt)