Shrook21/Code-Assistant

1from langchain_core.messages import HumanMessage, SystemMessage
2from agents.state import StateAgent
3from prompts.prompts import (
4    classify_prompt,
5    generate_prompt,
6    explain_prompt,
7    fallback_prompt
8)
9from tools.tools import retriever
10from langchain_community.llms import Ollama
11from config.settings import OLLAMA_MODEL_NAME
12import re
13
14llm = Ollama(
15    model=OLLAMA_MODEL_NAME,  # We can also use 'deepseek-coder:6.7b' or 'llama2:7b'
16    temperature=0.2
17)
18def chat(state: StateAgent) -> StateAgent:
19    user_input = state['message'][-1].content
20    prompt = classify_prompt(user_input)
21    result = llm.invoke(prompt)
22    raw = result.strip().lower()
23    match = re.search(r"(generate|explain|unclear)", raw)
24    task = match.group(1) if match else 'unclear'
25    return {**state, 'task': task, 'classification': raw}
26
27def router(state: StateAgent) -> str:
28    return state['task']
29
30def generate_code(state: StateAgent) -> StateAgent:
31    user_input = state['message'][-1].content
32    context = retriever(user_input)
33    prompt = generate_prompt(user_input, context)
34    output = llm.invoke(prompt)
35    return {
36        **state,
37        "message": state["message"] + [HumanMessage(content=prompt), SystemMessage(content=output)]
38    }
39
40def explain_code(state: StateAgent) -> StateAgent:
41    user_input = state['message'][-1].content
42    if not any(k in user_input.lower() for k in ['def ', 'class ', 'import ', 'for ', 'if ', 'while ', '=', 'print', 'return']):
43        output = f"I don't see any code in your input: '{user_input}'. Please provide the Python code you'd like me to explain."
44        return {**state, "message": state["message"] + [SystemMessage(content=output)]}
45    prompt = explain_prompt(user_input)
46    output = llm.invoke(prompt)
47    return {
48        **state,
49        "message": state["message"] + [HumanMessage(content=prompt), SystemMessage(content=output)]
50    }
51
52def fallback(state: StateAgent) -> StateAgent:
53    user_input = state['message'][-1].content
54    prompt = fallback_prompt(user_input)
55    output = llm.invoke(prompt)
56    return {
57        **state,
58        "message": state["message"] + [SystemMessage(content=output)]
59    }
60