|
1 | 1 | """ |
2 | 2 | Retrieval service |
3 | | -Handles query processing and retrieval chain operations |
| 3 | +Handles query processing and retrieval operations |
4 | 4 | """ |
5 | 5 |
|
6 | 6 | import logging |
7 | | -from langchain_openai import ChatOpenAI |
8 | 7 | from langchain_community.vectorstores import FAISS |
9 | | -from langchain.chains.retrieval import create_retrieval_chain |
10 | | -from langchain.chains.combine_documents import create_stuff_documents_chain |
11 | | -from langchain import hub |
12 | | -from langchain_core.language_models.chat_models import BaseChatModel |
13 | | -from langchain_core.language_models.llms import LLM |
14 | | -from langchain_core.outputs import LLMResult, Generation |
15 | | -from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, BaseMessage |
16 | | -from typing import List, Optional, Any |
17 | 8 | import config |
18 | 9 |
|
19 | 10 | logger = logging.getLogger(__name__) |
20 | 11 |
|
21 | 12 |
|
22 | | -class CustomLLM(LLM): |
23 | | - """ |
24 | | - Custom LLM class that uses the Llama-3.1-8B-Instruct endpoint |
25 | | - """ |
26 | | - |
27 | | - @property |
28 | | - def _llm_type(self) -> str: |
29 | | - """Return type of LLM.""" |
30 | | - return "custom_llm" |
31 | | - |
32 | | - def _call( |
33 | | - self, |
34 | | - prompt: str, |
35 | | - stop: Optional[List[str]] = None, |
36 | | - run_manager: Optional[Any] = None, |
37 | | - **kwargs: Any, |
38 | | - ) -> str: |
39 | | - """Call the LLM on the given prompt.""" |
40 | | - from .api_client import get_api_client |
41 | | - api_client = get_api_client() |
42 | | - return api_client.complete(prompt, max_tokens=kwargs.get('max_tokens', 150), temperature=kwargs.get('temperature', 0)) |
43 | | - |
44 | | - |
45 | | -class CustomChatModel(BaseChatModel): |
46 | | - """ |
47 | | - Custom Chat Model that uses the Llama-3.1-8B-Instruct endpoint |
48 | | - """ |
49 | | - |
50 | | - @property |
51 | | - def _llm_type(self) -> str: |
52 | | - """Return type of LLM.""" |
53 | | - return "custom_chat" |
54 | | - |
55 | | - def _generate( |
56 | | - self, |
57 | | - messages: List[BaseMessage], |
58 | | - stop: Optional[List[str]] = None, |
59 | | - run_manager: Optional[Any] = None, |
60 | | - **kwargs: Any, |
61 | | - ) -> LLMResult: |
62 | | - """Generate response from messages.""" |
63 | | - from .api_client import get_api_client |
64 | | - api_client = get_api_client() |
65 | | - |
66 | | - # Convert messages to a prompt string |
67 | | - # Build the prompt from all messages |
68 | | - prompt_parts = [] |
69 | | - |
70 | | - for msg in messages: |
71 | | - if isinstance(msg, SystemMessage): |
72 | | - prompt_parts.append(f"System: {msg.content}") |
73 | | - elif isinstance(msg, HumanMessage): |
74 | | - prompt_parts.append(f"User: {msg.content}") |
75 | | - elif isinstance(msg, AIMessage): |
76 | | - prompt_parts.append(f"Assistant: {msg.content}") |
77 | | - |
78 | | - # Join all parts and add assistant prompt suffix |
79 | | - full_prompt = "\n\n".join(prompt_parts) |
80 | | - if not full_prompt.endswith("Assistant:"): |
81 | | - full_prompt += "\n\nAssistant:" |
82 | | - |
83 | | - logger.info(f"Sending prompt to LLM (length: {len(full_prompt)} chars)") |
84 | | - |
85 | | - # Use the complete method which directly sends the prompt |
86 | | - # This calls: Llama-3.1-8B-Instruct/v1/completions with prompt |
87 | | - response_text = api_client.complete( |
88 | | - full_prompt, |
89 | | - max_tokens=kwargs.get('max_tokens', 150), |
90 | | - temperature=kwargs.get('temperature', 0) |
91 | | - ) |
92 | | - |
93 | | - generations = [Generation(text=response_text)] |
94 | | - return LLMResult(generations=[generations]) |
95 | | - |
96 | | - |
97 | | -def get_llm(api_key: str) -> BaseChatModel: |
98 | | - """ |
99 | | - Get LLM instance (ChatOpenAI or CustomChatModel based on config) |
100 | | -
|
101 | | - Args: |
102 | | - api_key: API key |
103 | | -
|
104 | | - Returns: |
105 | | - LLM instance |
106 | | - """ |
107 | | - # Check if using custom inference endpoint |
108 | | - if hasattr(config, 'INFERENCE_API_TOKEN') and config.INFERENCE_API_TOKEN: |
109 | | - return CustomChatModel() |
110 | | - else: |
111 | | - # Fallback to OpenAI ChatOpenAI |
112 | | - return ChatOpenAI( |
113 | | - model="gpt-3.5-turbo", |
114 | | - temperature=0, |
115 | | - openai_api_key=api_key |
116 | | - ) |
117 | | - |
118 | | - |
119 | | -def build_retrieval_chain(vectorstore: FAISS, api_key: str): |
120 | | - """ |
121 | | - Build retrieval chain with LLM (ChatOpenAI or CustomChatModel) |
122 | | - |
123 | | - Args: |
124 | | - vectorstore: FAISS vectorstore instance |
125 | | - api_key: API key |
126 | | - |
127 | | - Returns: |
128 | | - Configured retrieval chain |
129 | | - |
130 | | - Raises: |
131 | | - Exception: If chain building fails |
132 | | - """ |
133 | | - try: |
134 | | - retrieval_qa_chat_prompt = hub.pull("langchain-ai/retrieval-qa-chat") |
135 | | - llm = get_llm(api_key) |
136 | | - combine_docs_chain = create_stuff_documents_chain(llm, retrieval_qa_chat_prompt) |
137 | | - retrieval_chain = create_retrieval_chain( |
138 | | - vectorstore.as_retriever(search_kwargs={"k": 4}), |
139 | | - combine_docs_chain |
140 | | - ) |
141 | | - return retrieval_chain |
142 | | - except Exception as e: |
143 | | - logger.error(f"Error building retrieval chain: {str(e)}") |
144 | | - raise |
145 | | - |
146 | | - |
147 | 13 | def query_documents(query: str, vectorstore: FAISS, api_key: str) -> dict: |
148 | 14 | """ |
149 | 15 | Query the documents using RAG with custom embedding and inference |
|
0 commit comments