App.py 2.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  1. from typing import Optional
  2. import openai
  3. from embedchain.config import AppConfig, ChatConfig
  4. from embedchain.embedchain import EmbedChain
  5. class App(EmbedChain):
  6. """
  7. The EmbedChain app.
  8. Has two functions: add and query.
  9. adds(data_type, url): adds the data from the given URL to the vector db.
  10. query(query): finds answer to the given query using vector database and LLM.
  11. dry_run(query): test your prompt without consuming tokens.
  12. """
  13. def __init__(self, config: AppConfig = None, system_prompt: Optional[str] = None):
  14. """
  15. :param config: AppConfig instance to load as configuration. Optional.
  16. :param system_prompt: System prompt string. Optional.
  17. """
  18. if config is None:
  19. config = AppConfig()
  20. super().__init__(config, system_prompt)
  21. def get_llm_model_answer(self, prompt, config: ChatConfig):
  22. messages = []
  23. system_prompt = (
  24. self.system_prompt
  25. if self.system_prompt is not None
  26. else config.system_prompt
  27. if config.system_prompt is not None
  28. else None
  29. )
  30. if system_prompt:
  31. messages.append({"role": "system", "content": system_prompt})
  32. messages.append({"role": "user", "content": prompt})
  33. response = openai.ChatCompletion.create(
  34. model=config.model or "gpt-3.5-turbo-0613",
  35. messages=messages,
  36. temperature=config.temperature,
  37. max_tokens=config.max_tokens,
  38. top_p=config.top_p,
  39. stream=config.stream,
  40. )
  41. if config.stream:
  42. return self._stream_llm_model_response(response)
  43. else:
  44. return response["choices"][0]["message"]["content"]
  45. def _stream_llm_model_response(self, response):
  46. """
  47. This is a generator for streaming response from the OpenAI completions API
  48. """
  49. for line in response:
  50. chunk = line["choices"][0].get("delta", {}).get("content", "")
  51. yield chunk