@@ -4,16 +4,18 @@ class Chat:
44 def __init__ (self ):
55 self .client = None
66
7- def run (self , api_key = None , model = "sonar-medium-chat" , prompt = None , system_prompt = None , stream = None , max_tokens = None , temperature = None , top_p = None , top_k = None , presence_penalty = None , frequency_penalty = None ):
7+ def run (self , api_key = None , model = None , prompt = None , system_prompt = None , stream = None , max_tokens = None , temperature = None , top_p = None , top_k = None , presence_penalty = None , frequency_penalty = None ):
88
99 self .client = Client (api_key = api_key )
10- self .model = model if model else self .client .config .get ('chat_model ' )
10+ self .model = model if model else self .client .config .get ('base_chat_model ' )
1111
1212 conversation_history = []
1313
1414 if system_prompt :
1515 conversation_history .append ({"role" : "system" , "content" : system_prompt })
1616
17+ print ("Type 'exit' or 'quit' at any time to end the conversation.\n " )
18+
1719 print ("Assistant: Hello! How can I assist you today?" )
1820 while True :
1921 if prompt :
@@ -33,7 +35,7 @@ def run(self, api_key=None, model="sonar-medium-chat", prompt=None, system_promp
3335 conversation_history .append ({"role" : "user" , "content" : user_input })
3436
3537 payload = {
36- "model" : model ,
38+ "model" : self . model ,
3739 "messages" : conversation_history ,
3840 "system_prompt" : system_prompt ,
3941 "stream" : stream ,
@@ -63,10 +65,14 @@ class Search:
6365 def __init__ (self ):
6466 self .client = None
6567
66- def run (self , api_key = None , model = "sonar-medium-online" , query = None , system_prompt = None , stream = None , max_tokens = None , temperature = None , top_p = None , top_k = None , presence_penalty = None , frequency_penalty = None ):
68+ def run (self , api_key = None , model = None , query = None , system_prompt = None , stream = None , max_tokens = None , temperature = None , top_p = None , top_k = None , presence_penalty = None , frequency_penalty = None ):
6769
6870 self .client = Client (api_key = api_key )
69- self .model = model if model else self .client .config .get ('search_model' )
71+ self .model = model if model else self .client .config .get ('base_search_model' )
72+
73+ if "online" not in self .model :
74+ print ("Error: { Invalid model type }. Please use a search model instead of a chat model." )
75+ exit (1 )
7076
7177 if not query :
7278 print ("Error: { Invalid input detected }. Please enter a valid search query." )
@@ -81,7 +87,7 @@ def run(self, api_key=None, model="sonar-medium-online", query=None, system_prom
8187 message = [{"role" : "user" , "content" : query }]
8288
8389 payload = {
84- "model" : model ,
90+ "model" : self . model ,
8591 "messages" : message ,
8692 "system_prompt" : system_prompt ,
8793 "stream" : stream ,
@@ -103,4 +109,6 @@ def run(self, api_key=None, model="sonar-medium-online", query=None, system_prom
103109 else :
104110 response = self .client .post (endpoint , data )
105111 assistant_response = response
106- print (f"Assistant: { assistant_response } " )
112+ print (f"Assistant: { assistant_response } " )
113+
114+ print ("\n Thank you for using the Perplexity AI toolkit. Have a great day!" )
0 commit comments