refine admin initialization (#75)

rag/llm/chat_model.py

@@ -15,7 +15,7 @@
 #
 from abc import ABC
 from openai import OpenAI
-import os
+import openai
 
 
 class Base(ABC):
@@ -33,11 +33,14 @@ class GptTurbo(Base):
 
     def chat(self, system, history, gen_conf):
         if system: history.insert(0, {"role": "system", "content": system})
-        res = self.client.chat.completions.create(
-            model=self.model_name,
-            messages=history,
-            **gen_conf)
-        return res.choices[0].message.content.strip(), res.usage.completion_tokens
+        try:
+            res = self.client.chat.completions.create(
+                model=self.model_name,
+                messages=history,
+                **gen_conf)
+            return res.choices[0].message.content.strip(), res.usage.completion_tokens
+        except openai.APIError as e:
+            return "ERROR: "+str(e), 0
 
 
 from dashscope import Generation
@@ -58,7 +61,7 @@ class QWenChat(Base):
         )
         if response.status_code == HTTPStatus.OK:
             return response.output.choices[0]['message']['content'], response.usage.output_tokens
-        return response.message, 0
+        return "ERROR: " + response.message, 0
 
 
 from zhipuai import ZhipuAI
@@ -77,4 +80,4 @@ class ZhipuChat(Base):
         )
         if response.status_code == HTTPStatus.OK:
             return response.output.choices[0]['message']['content'], response.usage.completion_tokens
-        return response.message, 0
+        return "ERROR: " + response.message, 0

rag/nlp/__init__.py

@@ -1,7 +1,4 @@
-from . import search
-from rag.utils import ELASTICSEARCH
 
-retrievaler = search.Dealer(ELASTICSEARCH)
 
 from nltk.stem import PorterStemmer
 stemmer = PorterStemmer()
@@ -39,10 +36,12 @@ BULLET_PATTERN = [[
 ]
 ]
 
+
 def random_choices(arr, k):
     k = min(len(arr), k)
     return random.choices(arr, k=k)
 
+
 def bullets_category(sections):
     global BULLET_PATTERN
     hits = [0] * len(BULLET_PATTERN)

rag/nlp/search.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 import json
 import re
-from elasticsearch_dsl import Q, Search, A
+from elasticsearch_dsl import Q, Search
 from typing import List, Optional, Dict, Union
 from dataclasses import dataclass
 
@@ -183,6 +183,7 @@ class Dealer:
 
     def insert_citations(self, answer, chunks, chunk_v,
                          embd_mdl, tkweight=0.3, vtweight=0.7):
+        assert len(chunks) == len(chunk_v)
         pieces = re.split(r"([；。？!！\n]|[a-z][.?;!][ \n])", answer)
         for i in range(1, len(pieces)):
             if re.match(r"[a-z][.?;!][ \n]", pieces[i]):
@@ -216,7 +217,7 @@ class Dealer:
             if mx < 0.55:
                 continue
             cites[idx[i]] = list(
-                set([str(i) for i in range(len(chunk_v)) if sim[i] > mx]))[:4]
+                set([str(ii) for ii in range(len(chunk_v)) if sim[ii] > mx]))[:4]
 
         res = ""
         for i, p in enumerate(pieces):
@@ -225,6 +226,7 @@ class Dealer:
                 continue
             if i not in cites:
                 continue
+            assert int(cites[i]) < len(chunk_v)
             res += "##%s$$" % "$".join(cites[i])
 
         return res