Update main.py

main.py

@@ -1,16 +1,29 @@
-#!/usr/bin/env python3
-"""
-Solution for task 6997111cd6d3a5544a3deffd.
-This file contains a minimal example implementation.
-"""
+import os
+from langchain_openai import ChatOpenAI
 
-def greet(name: str) -> str:
-    """Return a greeting message for the given name."""
-    return f"Hello, {name}!"
+# Configure LLM to connect to local LM Studio server
+llm = ChatOpenAI(
+    model="gpt-4o-mini",  # or any model name supported by the local server
+    temperature=0,
+    base_url="http://localhost:1234/v1",
+    api_key="lm-studio"
+)
+
+def ask(prompt: str) -> str:
+    """Send a prompt to the LLM and return its response."""
+    result = llm.invoke({"input": prompt})
+    # The output is in result.content
+    return result.content
 
 if __name__ == "__main__":
-    import sys
-    if len(sys.argv) > 1:
-        print(greet(sys.argv[1]))
-    else:
-        print("Usage: python main.py <name>")
+    print("Simple AI Agent – type your question (Ctrl+C to exit).")
+    while True:
+        try:
+            user_input = input("\n> ")
+            if not user_input.strip():
+                continue
+            response = ask(user_input)
+            print(f"\nAnswer: {response}")
+        except KeyboardInterrupt:
+            print("\nExiting.")
+            break