File size: 3,939 Bytes
a0c836f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | from langchain_groq import ChatGroq
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from state import State
from langchain_core.messages import HumanMessage, SystemMessage
import json
import base64
from langchain_google_genai import ChatGoogleGenerativeAI
from prompt import *
from dotenv import load_dotenv
load_dotenv()
llm = ChatGroq(model="llama-3.1-8b-instant", temperature=0.6)
vision_llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash-lite", temperature=0.0)
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vdb = Chroma(
embedding_function=embeddings,
persist_directory="./chroma_db"
)
def retriever_agent(state: State) -> str:
rewritten_query = state.get("rewritten_query")
retriever = vdb.as_retriever(search_kwargs={"k": 6})
results = retriever.invoke(rewritten_query)
content = "\n\n".join([doc.page_content for doc in results])
return {
"content": content
}
def rewritten_query_agent(state: State) -> str:
user_input = state.get("query")
chat_history = state.get("messages")
messages = [
SystemMessage(content=REWRITE_PROMPT),
HumanMessage(content=query_rewrite_extend(user_input, chat_history))
]
try:
response = llm.invoke(messages)
rewritten_query = response.content
return {
"rewritten_query": rewritten_query
}
except Exception as e:
print(f"Error in rewritten_query: {e}")
return None
def response_agent(state: State) -> str:
# Use the user's ORIGINAL wording for the reply (so the answer matches their
# language, e.g. Arabic), while retrieval already used the rewritten query.
user_query = state.get("query") or state.get("rewritten_query")
chat_history = state.get("messages")
content = state.get("content")
messages = [
SystemMessage(content=SYSTEM_PROMPT),
HumanMessage(content=system_prompt_extend(user_query, chat_history, content))
]
try:
response = llm.invoke(messages)
answer = response.content
return {
"response": answer }
except Exception as e:
print(f"Error in response_agent: {e}")
return None
def check_input_type(state: State) -> str:
if state.get("image"):
return "image"
return "text"
def vision_agent(state: State) -> str:
image_bytes = state.get("image")
user_query = state.get("query")
image_b64 = base64.b64encode(image_bytes).decode("utf-8")
message = HumanMessage(
content=[
{"type": "text", "text": VISION_PROMPT},
{"type": "image_url", "image_url": f"data:image/jpeg;base64,{image_b64}"}
]
)
try:
response = vision_llm.invoke([message])
raw_text = response.content.strip()
if raw_text.startswith("```"):
raw_text = raw_text.strip("`").replace("json", "", 1).strip()
parsed = json.loads(raw_text)
medicine_name = parsed.get("medicine_name")
if not medicine_name:
return {"rewritten_query": None, "content": "NOT_FOUND"}
if user_query and user_query.strip():
rewritten_query = f"{user_query.strip()} for {medicine_name}"
else:
rewritten_query = f"Tell me about {medicine_name}"
return {"rewritten_query": rewritten_query}
except Exception as e:
print(f"Error in vision_agent: {e}")
return {"rewritten_query": None, "content": "NOT_FOUND"}
def generate_title(first_message: str) -> str:
messages = [
SystemMessage(content=TITLE_PROMPT),
HumanMessage(content=first_message)
]
try:
response = llm.invoke(messages)
return response.content.strip()[:100]
except Exception as e:
print(f"Error in generate_title: {e}")
return "Medical Consultation" |