Upload 2 files

app.py

@@ -0,0 +1,163 @@
+import os
+import pickle
+
+import streamlit as st
+from dotenv import load_dotenv
+from laas import ChatLaaS
+from langchain.embeddings import CacheBackedEmbeddings
+from langchain.retrievers import ContextualCompressionRetriever, EnsembleRetriever
+from langchain.retrievers.document_compressors import (
+    CrossEncoderReranker,
+    FlashrankRerank,
+)
+from langchain_core.vectorstores import VectorStore
+from langchain.storage import LocalFileStore
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+from langchain_community.document_loaders.generic import GenericLoader
+from langchain_community.document_loaders.parsers.language.language_parser import (
+    LanguageParser,
+)
+from langchain_community.retrievers import BM25Retriever
+from langchain_community.vectorstores import FAISS
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnableLambda, RunnablePassthrough
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_text_splitters import Language, RecursiveCharacterTextSplitter
+
+# Load environment variables
+load_dotenv()
+
+# Set up environment variables
+os.environ["LANGCHAIN_TRACING_V2"] = "true"
+os.environ["LANGCHAIN_PROJECT"] = "Code QA Bot"
+
+
+@st.cache_resource
+def setup_embeddings_and_db(project_folder: str):  # Note the underscore before 'docs'
+    CACHE_ROOT_PATH = os.path.join(os.path.expanduser("~"), ".cache")
+    CACHE_MODELS_PATH = os.path.join(CACHE_ROOT_PATH, "models")
+    CACHE_EMBEDDINGS_PATH = os.path.join(CACHE_ROOT_PATH, "embeddings")
+
+    if not os.path.exists(CACHE_MODELS_PATH):
+        os.makedirs(CACHE_MODELS_PATH)
+    if not os.path.exists(CACHE_EMBEDDINGS_PATH):
+        os.makedirs(CACHE_EMBEDDINGS_PATH)
+
+    store = LocalFileStore(CACHE_EMBEDDINGS_PATH)
+
+    model_name = "BAAI/bge-m3"
+    model_kwargs = {"device": "mps"}
+    encode_kwargs = {"normalize_embeddings": False}
+    embeddings = HuggingFaceEmbeddings(
+        model_name=model_name,
+        model_kwargs=model_kwargs,
+        encode_kwargs=encode_kwargs,
+        cache_folder=CACHE_MODELS_PATH,
+        multi_process=False,
+        show_progress=True,
+    )
+
+    cached_embeddings = CacheBackedEmbeddings.from_bytes_store(
+        embeddings,
+        store,
+        namespace=embeddings.model_name,
+    )
+
+    FAISS_DB_INDEX = os.path.join(project_folder, "langchain_faiss")
+    db = FAISS.load_local(
+        FAISS_DB_INDEX,  # 로드할 FAISS 인덱스의 디렉토리 이름
+        cached_embeddings,  # 임베딩 정보를 제공
+        allow_dangerous_deserialization=True,  # 역직렬화를 허용하는 옵션
+    )
+
+    return db
+
+
+# Function to set up retrievers and chain
+@st.cache_resource
+def setup_retrievers_and_chain(
+    _db: VectorStore, project_folder: str
+):  # Note the underscores
+    faiss_retriever = _db.as_retriever(search_type="mmr", search_kwargs={"k": 20})
+
+    bm25_retriever_path = os.path.join(project_folder, "bm25_retriever.pkl")
+    with open(bm25_retriever_path, "rb") as f:
+        bm25_retriever = pickle.load(f)
+        bm25_retriever.k = 20
+
+    ensemble_retriever = EnsembleRetriever(
+        retrievers=[bm25_retriever, faiss_retriever],
+        weights=[0.6, 0.4],
+        search_type="mmr",
+    )
+
+    model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-v2-m3")
+    compressor = CrossEncoderReranker(model=model, top_n=5)
+    compression_retriever = ContextualCompressionRetriever(
+        base_compressor=compressor,
+        base_retriever=ensemble_retriever,
+    )
+
+    laas = ChatLaaS(
+        project=st.secrets["LAAS_PROJECT"],
+        api_key=st.secrets["LAAS_API_KEY"],
+        hash=st.secrets["LAAS_HASH"],
+    )
+
+    rag_chain = (
+        {
+            "context": compression_retriever | RunnableLambda(lambda x: str(x)),
+            "question": RunnablePassthrough(),
+        }
+        | RunnableLambda(
+            lambda x: laas.invoke(
+                "", params={"context": x["context"], "question": x["question"]}
+            )
+        )
+        | StrOutputParser()
+    )
+
+    return rag_chain
+
+
+def main():
+    st.title("Code QA Bot")
+
+    # Initialize session state for project folder and answer
+    if "project_folder" not in st.session_state:
+        st.session_state.project_folder = ""
+    if "answer" not in st.session_state:
+        st.session_state.answer = ""
+
+    # 프로젝트 경로 입력 받기
+    project_folder = st.text_input(
+        "Enter the project folder path:", value=st.session_state.project_folder
+    )
+    st.session_state.project_folder = project_folder
+
+    if project_folder:
+        # 프로젝트 경로가 입력되면 벡터 스토어와 체인 설정
+        db = setup_embeddings_and_db(project_folder)
+        rag_chain = setup_retrievers_and_chain(db, project_folder)
+
+        # 사용자 질문 입력 받기
+        user_question = st.text_input("Ask a question about the code:")
+
+        # Add a button to reset the answer
+        if st.button("Reset Answer"):
+            st.session_state.answer = ""
+
+        if user_question:
+            with st.spinner("Generating answer..."):
+                response = rag_chain.invoke(user_question)
+                st.session_state.answer = response
+
+        # Display the answer
+        if st.session_state.answer:
+            st.write(st.session_state.answer)
+    else:
+        st.warning("Please enter the project folder path to proceed.")
+
+
+if __name__ == "__main__":
+    main()

laas.py

@@ -0,0 +1,80 @@
+import logging
+from typing import Any, List, Optional
+
+import requests
+from langchain_core.callbacks import CallbackManagerForLLMRun
+from langchain_core.language_models import BaseChatModel, BaseLanguageModel
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
+from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
+from langchain_core.pydantic_v1 import BaseModel, Field, SecretStr, root_validator
+
+logger = logging.getLogger(__name__)
+
+
+class ChatLaaS(BaseChatModel):
+    laas_api_base: Optional[str] = Field(
+        default="https://api-laas.wanted.co.kr/api/preset", alias="base_url"
+    )
+    laas_api_key: Optional[SecretStr] = Field(default=None, alias="api_key")
+    laas_project: Optional[str] = Field(default=None, alias="project")
+    laas_hash: Optional[str] = Field(default=None, alias="hash")
+    timeout: Optional[float] = Field(default=60.0)
+
+    _ROLE_MAP = {
+        "human": "user",
+        "ai": "assistant",
+    }
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of chat model."""
+        return "laas-chat"
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """Return whether this model can be serialized by Langchain."""
+        return False
+
+    def _generate(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        try:
+            body = {
+                "hash": self.laas_hash,
+                "messages": [
+                    {
+                        "role": self._ROLE_MAP.get(msg.type, msg.type),
+                        "content": msg.content,
+                    }
+                    for msg in messages
+                    if msg.content.strip()  # This filters out empty or whitespace-only content
+                ],
+                **kwargs,
+            }
+            print(body)
+            # return
+
+            headers = {
+                "Content-Type": "application/json",
+                "apiKey": self.laas_api_key.get_secret_value(),
+                "project": self.laas_project,
+            }
+
+            response = requests.post(
+                f"{self.laas_api_base}/chat/completions",
+                headers=headers,
+                json=body,
+                timeout=self.timeout,
+            ).json()
+
+            # Extract the content from the API response
+            content = response["choices"][0]["message"]["content"]
+            message = AIMessage(id=response["id"], content=content)
+            generation = ChatGeneration(message=message)
+            return ChatResult(generations=[generation])
+        except Exception as e:
+            raise