Spaces:
Paused
Paused
| import requests | |
| from langchain.llms import BaseLLM # type: ignore | |
| from langchain_core.callbacks import CallbackManagerForLLMRun | |
| from langchain_core.outputs.llm_result import LLMResult | |
| from typing import List, Optional, Any | |
| from openai import OpenAI | |
| import os | |
| api_key_from_env = os.getenv("API_KEY_PERPLEXITY") | |
| class PerplexityCrewLLM(BaseLLM): | |
| api_key: str | |
| model_name: str | |
| def call_perplexity_ai(self, prompt: str) -> LLMResult: | |
| url = "https://api.perplexity.ai/chat/completions" | |
| payload = { | |
| "model": self.model_name, | |
| "messages": [ | |
| { | |
| "role": "system", | |
| "content": "Be precise and concise." | |
| }, | |
| { | |
| "role": "user", | |
| "content": prompt | |
| } | |
| ] | |
| } | |
| headers = { | |
| "Authorization": f"Bearer {self.api_key}", | |
| "accept": "application/json", | |
| "content-type": "application/json" | |
| } | |
| response = requests.post(url, json=payload, headers=headers) | |
| # Convert the response JSON to dictionary | |
| json_response = response.json() | |
| return json_response | |
| def _generate(self, prompts: List[str], stop: Optional[List[str]] = None, | |
| run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any) -> LLMResult: | |
| generations = [] | |
| for prompt in prompts: | |
| generations.append([self._call(prompt, stop=stop, **kwargs)]) | |
| return LLMResult.construct(generations=generations) | |
| def _call(self, prompt: str, stop: Optional[List[str]] = None, max_tokens: Optional[int] = None) -> LLMResult: | |
| response_data = self.call_perplexity_ai(prompt) | |
| model = LLMResult.construct(text=response_data['choices'][0]["message"]["content"]) # type: ignore | |
| return model | |
| def _llm_type(self) -> str: | |
| return "PerplexityAI" | |
| def PerplexitySearchLLM(api_key,model_name="sonar-medium-online",base_url="https://api.perplexity.ai"): | |
| client = OpenAI(api_key=api_key_from_env, base_url=base_url) | |
| def call_model(query:str): | |
| messages = [ | |
| #It is recommended to use only single-turn conversations and avoid system prompts for the online LLMs (sonar-small-online and sonar-medium-online). | |
| # { | |
| # "role": "system", | |
| # "content": ( | |
| # "You are an artificial intelligence assistant and you need to " | |
| # "engage in a helpful, detailed, polite conversation with a user." | |
| # ), | |
| # }, | |
| { | |
| "role": "user", | |
| "content": ( | |
| query | |
| ), | |
| }, | |
| ] | |
| response = client.chat.completions.create( | |
| model=model_name, | |
| messages=messages, # type: ignore | |
| ) | |
| result = response.choices[0].message.content #only the text is returned | |
| return result | |
| return call_model | |
| call_llm = PerplexitySearchLLM(api_key=api_key_from_env,model_name="sonar-medium-online") | |
| def perplexity_search(search_query: str): | |
| return call_llm(search_query) |