from sentence_transformers import SentenceTransformer
from together import Together
import vecs
import os
import gradio as gr

user = os.getenv("user")
password = os.getenv("password")
host = os.getenv("host")
port = os.getenv("port")
db_name = "postgres"
DB_CONNECTION = f"postgresql://{user}:{password}@{host}:{port}/{db_name}"
vx = vecs.create_client(DB_CONNECTION)
model = SentenceTransformer('Snowflake/snowflake-arctic-embed-xs')
client = Together(api_key=os.getenv('TOGETHER_API_KEY'))

login_user = os.getenv('login_user')
login_pass = os.getenv('login_pass')


class Ticio():

  def __init__(self, system_prompt, model = 'meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo'):
    self.model = model
    self.messages = [{'role': 'system', 'content': system_prompt}]

  def inference(self):
    response = client.chat.completions.create(
        model=self.model,
        messages=self.messages,
    )
    self.messages.append({"role": 'assistant', "content": response.choices[0].message.content})
    return response

  def add_message(self, content, role):
    self.messages.append({"role": role, "content": content})
    return self.messages

  def last_inference(self):
    return self.messages[-1]['content']

system_prompt = '''Eres Ticio, un asistente de investigación jurídica. Tu objetivo es reponder a las preguntas planteadas por el usuario, sistematizando la información de contexto que tengas. Para esto tienes las siguiente funciones:

#usa esta función cuando te hagan preguntas que no puedas responder con la información disponible en el contexto, el query es la pregunta del usuario.
def buscar(query):
  vector = model.encode(query)
  información = query_db(vector)
  return informaciòn

cuando quieras usar una función, responde SOLO con la siguiente sintaxis:
<tool_call>function(parameters)</tool_call>

Siempre debes buscar información que no puedas responder directamente con la información de contexto que tengas en el chat, si no tienes información siempre debes buscar. Cuando tengas la información debes responder citando diractamente la información
Debes presentar información resumida que ayude a la persona a identificar de forma rapida si un documento le es relevante. SOLO HAZ ESTO DESPUES DE HABER BUSCADO INFORMACIÓN
'''
def function_call(calls):
  results = []
  for item in calls:
    a = eval(item)
    results.append(a)
  return results

def parse_call(output):
  token = '<tool_call>'
  token2= '</tool_call>'

  calls = []
  if token in output:
    split = output.split(token)
    for item in split:
      if token2 in item:
        calls.append(item.split(token2)[0])
  return calls

def query_db(query, limit = 5, filters = {}, measure = "cosine_distance", include_value = True, include_metadata=True, table = "CE"):
  query_embeds = vx.get_or_create_collection(name= table, dimension=384)
  ans = query_embeds.query(
      data=query,
      limit=limit,
      filters=filters,
      measure=measure,
      include_value=include_value,
      include_metadata=include_metadata,
  )
  return ans

def sort_by_score(item):
  return item[1]

def construct_result(ans):
  ans.sort(key=sort_by_score, reverse=True)
  results = ""
  for i in range(0, len(ans)):
    a, b = ans[i][2].get("documento"), ans[i][2].get("fragmento")
    results += (f"En el documento {a}, se dijo {b}\n")
  return results


def buscar(query):
  results = query_db(model.encode(query))
  return construct_result(results)


def inference(message, history):
  ticio = Ticio(system_prompt)
  ticio.add_message(message, 'user')
  for item in history:
    content = item.get('content')
    role = item.get('role')
    ticio.add_message(content, role)
  ticio.inference()
  response = ticio.last_inference()
  calls = parse_call(response)
  if calls:
    results = function_call(calls)
    for result in results:
      ticio.add_message(result, 'assistant')
    ticio.add_message(message, 'user')
    ticio.inference()
    return ticio.last_inference()
  else:
    return response

theme = gr.themes.Base(
    primary_hue="red",
    secondary_hue="red",
    neutral_hue="neutral",
).set(
    button_primary_background_fill='#910A0A',
    button_primary_border_color='*primary_300',
    button_primary_text_color='*primary_50'
)

demo = gr.ChatInterface(fn=inference, type="messages", examples=["¿Qué se ha dicho acerca de la expresión 'devengado' cuando la usa la UGPP?"], title="Ticio", theme=theme, cache_examples=False, api_name=False)
demo.launch(auth=(login_user, login_pass), share = True, ssr_mode=False)