initial commit

app.py

@@ -0,0 +1,134 @@
+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 message in history:
+    content = message.get('content')
+    role = message.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)
+demo.launch(auth=[login_user, login_pass])